test: refactor

tests/cli_kv_test.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+
+import base64
+import random
+import tempfile
+
+from config.node_config import GENESIS_ACCOUNT
+from utility.submission import ENTRY_SIZE, bytes_to_entries
+from utility.utils import (
+    assert_equal,
+    wait_until,
+)
+from utility.kv import to_stream_id
+from test_framework.test_framework import TestFramework
+
+
+class KVTest(TestFramework):
+    def setup_params(self):
+        self.num_blockchain_nodes = 1
+        self.num_nodes = 1
+
+    def run_test(self):
+        self.setup_kv_node(0, [to_stream_id(0)])
+        self._kv_write_use_cli(
+            self.blockchain_nodes[0].rpc_url,
+            self.contract.address(),
+            GENESIS_ACCOUNT.key,
+            self.nodes[0].rpc_url,
+            None,
+            to_stream_id(0),
+            "0,1,2,3,4,5,6,7,8,9,10",
+            "0,1,2,3,4,5,6,7,8,9,10"
+        )
+        wait_until(lambda: self.kv_nodes[0].kv_get_trasanction_result(0) == "Commit")
+        res = self._kv_read_use_cli(
+            self.kv_nodes[0].rpc_url,
+            to_stream_id(0),
+            "0,1,2,3,4,5,6,7,8,9,10,11"
+        )
+        for i in range(11):
+            assert_equal(res[str(i)], str(i))
+        assert_equal(res["11"], "")
+
+if __name__ == "__main__":
+    KVTest().main()

tests/skip_tx_test.py → tests/cli_skip_tx_test.py

@@ -17,10 +17,10 @@
 class SkipTxTest(TestFramework):
     def setup_params(self):
         self.num_blockchain_nodes = 1
-        self.num_nodes = 2
+        self.num_nodes = 1
 
     def run_test(self):
-        node_idx = random.randint(0, self.num_nodes - 1)
+        node_idx = 0 
 
         file_to_upload = tempfile.NamedTemporaryFile(dir=self.root_dir, delete=False)
         data = random.randbytes(256 * 2048) 

tests/test_framework/blockchain_node.py

@@ -34,6 +34,7 @@ def block_time(self):
 class NodeType(Enum):
     BlockChain = 0
     Zgs = 1
+    KV = 2
 
 
 class FailedToStartError(Exception):

tests/test_framework/test_framework.py

@@ -11,6 +11,7 @@
 import tempfile
 import time
 import traceback
+import json
 from pathlib import Path
 
 from eth_utils import encode_hex
@@ -390,6 +391,118 @@ def _download_file_use_cli(
         os.remove(file_to_download)
 
         return
+
+    def _kv_write_use_cli(
+        self,
+        blockchain_node_rpc_url,
+        contract_address,
+        key,
+        node_rpc_url,
+        indexer_url,
+        stream_id,
+        kv_keys,
+        kv_values,
+        skip_tx = True,
+    ):
+        kv_write_args = [
+            self.cli_binary,
+            "kv-write",
+            "--url",
+            blockchain_node_rpc_url,
+            "--contract",
+            contract_address,
+            "--key",
+            encode_hex(key),
+            "--skip-tx="+str(skip_tx),
+            "--stream-id",
+            stream_id,
+            "--stream-keys",
+            kv_keys,
+            "--stream-values",
+            kv_values,
+            "--log-level",
+            "debug",
+            "--gas-limit",
+            "10000000",
+        ]
+        if node_rpc_url is not None:
+            kv_write_args.append("--node")
+            kv_write_args.append(node_rpc_url)
+        elif indexer_url is not None:
+            kv_write_args.append("--indexer")
+            kv_write_args.append(indexer_url)
+        self.log.info("kv write with cli: {}".format(kv_write_args))
+
+        output = tempfile.NamedTemporaryFile(dir=self.root_dir, delete=False, prefix="zgs_client_output_")
+        output_name = output.name
+        output_fileno = output.fileno()
+
+        try:
+            proc = subprocess.Popen(
+                kv_write_args,
+                text=True,
+                stdout=output_fileno,
+                stderr=output_fileno,
+            )
+
+            return_code = proc.wait(timeout=60)
+
+            output.seek(0)
+            lines = output.readlines()
+        except Exception as ex:
+            self.log.error("Failed to write kv via CLI tool, output: %s", output_name)
+            raise ex
+        finally:
+            output.close()
+
+        assert return_code == 0, "%s write kv failed, output: %s, log: %s" % (self.cli_binary, output_name, lines)
+
+        return
+
+    def _kv_read_use_cli(
+        self,
+        node_rpc_url,
+        stream_id,
+        kv_keys
+    ):
+        kv_read_args = [
+            self.cli_binary,
+            "kv-read",
+            "--node",
+            node_rpc_url,
+            "--stream-id",
+            stream_id,
+            "--stream-keys",
+            kv_keys,
+            "--log-level",
+            "debug",
+        ]
+        self.log.info("kv read with cli: {}".format(kv_read_args))
+
+        output = tempfile.NamedTemporaryFile(dir=self.root_dir, delete=False, prefix="zgs_client_output_")
+        output_name = output.name
+        output_fileno = output.fileno()
+
+        try:
+            proc = subprocess.Popen(
+                kv_read_args,
+                text=True,
+                stdout=output_fileno,
+                stderr=output_fileno,
+            )
+
+            return_code = proc.wait(timeout=60)
+            output.seek(0)
+            lines = output.readlines()
+        except Exception as ex:
+            self.log.error("Failed to read kv via CLI tool, output: %s", output_name)
+            raise ex
+        finally:
+            output.close()
+
+        assert return_code == 0, "%s read kv failed, output: %s, log: %s" % (self.cli_binary, output_name, lines)
+
+        return json.loads(lines[0].decode("utf-8").strip())
 
     def setup_params(self):
         self.num_blockchain_nodes = 1

tests/utility/kv.py

@@ -0,0 +1,203 @@
+from enum import Enum
+import random
+
+
+class AccessControlOps(Enum):
+    GRANT_ADMIN_ROLE = 0x00
+    RENOUNCE_ADMIN_ROLE = 0x01
+    SET_KEY_TO_SPECIAL = 0x10
+    SET_KEY_TO_NORMAL = 0x11
+    GRANT_WRITER_ROLE = 0x20
+    REVOKE_WRITER_ROLE = 0x21
+    RENOUNCE_WRITER_ROLE = 0x22
+    GRANT_SPECIAL_WRITER_ROLE = 0x30
+    REVOKE_SPECIAL_WRITER_ROLE = 0x31
+    RENOUNCE_SPECIAL_WRITER_ROLE = 0x32
+
+    @staticmethod
+    def grant_admin_role(stream_id, address):
+        return [AccessControlOps.GRANT_ADMIN_ROLE, stream_id, to_address(address)]
+
+    @staticmethod
+    def renounce_admin_role(stream_id):
+        return [AccessControlOps.RENOUNCE_ADMIN_ROLE, stream_id]
+
+    @staticmethod
+    def set_key_to_special(stream_id, key):
+        return [AccessControlOps.SET_KEY_TO_SPECIAL, stream_id, key]
+
+    @staticmethod
+    def set_key_to_normal(stream_id, key):
+        return [AccessControlOps.SET_KEY_TO_NORMAL, stream_id, key]
+
+    @staticmethod
+    def grant_writer_role(stream_id, address):
+        return [AccessControlOps.GRANT_WRITER_ROLE, stream_id, to_address(address)]
+
+    @staticmethod
+    def revoke_writer_role(stream_id, address):
+        return [AccessControlOps.REVOKE_WRITER_ROLE, stream_id, to_address(address)]
+
+    @staticmethod
+    def renounce_writer_role(stream_id):
+        return [AccessControlOps.RENOUNCE_WRITER_ROLE, stream_id]
+
+    @staticmethod
+    def grant_special_writer_role(stream_id, key, address):
+        return [
+            AccessControlOps.GRANT_SPECIAL_WRITER_ROLE,
+            stream_id,
+            key,
+            to_address(address),
+        ]
+
+    @staticmethod
+    def revoke_special_writer_role(stream_id, key, address):
+        return [
+            AccessControlOps.REVOKE_SPECIAL_WRITER_ROLE,
+            stream_id,
+            key,
+            to_address(address),
+        ]
+
+    @staticmethod
+    def renounce_special_writer_role(stream_id, key):
+        return [AccessControlOps.RENOUNCE_SPECIAL_WRITER_ROLE, stream_id, key]
+
+
+op_with_key = [
+    AccessControlOps.SET_KEY_TO_SPECIAL,
+    AccessControlOps.SET_KEY_TO_NORMAL,
+    AccessControlOps.GRANT_SPECIAL_WRITER_ROLE,
+    AccessControlOps.REVOKE_SPECIAL_WRITER_ROLE,
+    AccessControlOps.RENOUNCE_SPECIAL_WRITER_ROLE,
+]
+
+op_with_address = [
+    AccessControlOps.GRANT_ADMIN_ROLE,
+    AccessControlOps.GRANT_WRITER_ROLE,
+    AccessControlOps.REVOKE_WRITER_ROLE,
+    AccessControlOps.GRANT_SPECIAL_WRITER_ROLE,
+    AccessControlOps.REVOKE_SPECIAL_WRITER_ROLE,
+]
+
+
+MAX_STREAM_ID = 100
+MAX_DATA_LENGTH = 256 * 1024 * 4
+MIN_DATA_LENGTH = 10
+MAX_U64 = (1 << 64) - 1
+MAX_KEY_LEN = 2000
+
+STREAM_DOMAIN = bytes.fromhex(
+    "df2ff3bb0af36c6384e6206552a4ed807f6f6a26e7d0aa6bff772ddc9d4307aa"
+)
+
+
+def with_prefix(x):
+    x = x.lower()
+    if not x.startswith("0x"):
+        x = "0x" + x
+    return x
+
+
+def pad(x, l):
+    ans = hex(x)[2:]
+    return "0" * (l - len(ans)) + ans
+
+
+def to_address(x):
+    if x.startswith("0x"):
+        return x[2:]
+    return x
+
+
+def to_stream_id(x):
+    return pad(x, 64)
+
+
+def to_key_with_size(x):
+    size = pad(len(x) // 2, 6)
+    return size + x
+
+
+def rand_key():
+    len = random.randrange(1, MAX_KEY_LEN)
+    if len % 2 == 1:
+        len += 1
+    return "".join([hex(random.randrange(16))[2:] for i in range(len)])
+
+
+def rand_write(stream_id=None, key=None, size=None):
+    return [
+        (
+            to_stream_id(random.randrange(0, MAX_STREAM_ID))
+            if stream_id is None
+            else stream_id
+        ),
+        rand_key() if key is None else key,
+        random.randrange(MIN_DATA_LENGTH, MAX_DATA_LENGTH) if size is None else size,
+    ]
+
+
+def is_access_control_permission_denied(x):
+    if x is None:
+        return False
+    return x.startswith("AccessControlPermissionDenied")
+
+
+def is_write_permission_denied(x):
+    if x is None:
+        return False
+    return x.startswith("WritePermissionDenied")
+
+
+# reads: array of [stream_id, key]
+# writes: array of [stream_id, key, data_length]
+
+
+def create_kv_data(version, reads, writes, access_controls):
+    # version
+    data = bytes.fromhex(pad(version, 16))
+    tags = []
+    # read set
+    data += bytes.fromhex(pad(len(reads), 8))
+    for read in reads:
+        data += bytes.fromhex(read[0])
+        data += bytes.fromhex(to_key_with_size(read[1]))
+    # write set
+    data += bytes.fromhex(pad(len(writes), 8))
+    # write set meta
+    for write in writes:
+        data += bytes.fromhex(write[0])
+        data += bytes.fromhex(to_key_with_size(write[1]))
+        data += bytes.fromhex(pad(write[2], 16))
+        tags.append(write[0])
+        if len(write) == 3:
+            write_data = random.randbytes(write[2])
+            write.append(write_data)
+    # write data
+    for write in writes:
+        data += write[3]
+    # access controls
+    data += bytes.fromhex(pad(len(access_controls), 8))
+    for ac in access_controls:
+        k = 0
+        # type
+        data += bytes.fromhex(pad(ac[k].value, 2))
+        k += 1
+        # stream_id
+        tags.append(ac[k])
+        data += bytes.fromhex(ac[k])
+        k += 1
+        # key
+        if ac[0] in op_with_key:
+            data += bytes.fromhex(to_key_with_size(ac[k]))
+            k += 1
+        # address
+        if ac[0] in op_with_address:
+            data += bytes.fromhex(ac[k])
+            k += 1
+    tags = list(set(tags))
+    tags = sorted(tags)
+    tags = STREAM_DOMAIN + bytes.fromhex("".join(tags))
+    return data, tags