fix stuff

Author: erwin-wee

chainbench/profile/ethereum/subscriptions.py

@@ -0,0 +1,16 @@
+from chainbench.user import EvmUser, WssJrpcUser
+from chainbench.user.protocol.ethereum import EthSubscribe
+
+
+class EthSubscriptions(WssJrpcUser, EvmUser):
+    subscriptions = [
+        EthSubscribe(["newHeads"]),
+        # logs subscription for approve method signature
+        EthSubscribe(["logs", {"topics": ["0x8c5be1e5ebec7d5bd14f71427d1e84f3dd0314c0f7b2291e5b200ac8c7c3b925"]}]),
+        EthSubscribe(["newPendingTransactions"]),
+    ]
+
+    def get_notification_name(self, parsed_response: dict):
+        return self.get_subscription(
+            subscription_id=parsed_response["params"]["subscription"]
+        ).subscribe_rpc_call.params[0]

chainbench/test_data/blockchain.py

@@ -47,7 +47,7 @@ class Block:
     block_number: BlockNumber
 
     def to_json(self) -> str:
-        return json.dumps(self.__dict__)
+        return json.dumps(self.__dict__).decode("utf-8")
 
 
 @dataclass

chainbench/user/__init__.py

@@ -3,7 +3,7 @@
 
 from .common import get_subclass_tasks
 from .http import HttpUser
-from .jsonrpc import JsonRpcUser
+from .jsonrpc import JrpcHttpUser
 from .wss import WssJrpcUser
 
 # importing plugins here as all profiles depend on it
@@ -15,7 +15,7 @@
     "EthBeaconUser",
     "EvmUser",
     "HttpUser",
-    "JsonRpcUser",
+    "JrpcHttpUser",
     "SolanaUser",
     "StarkNetUser",
     "WssJrpcUser",

chainbench/user/jsonrpc.py

@@ -12,7 +12,7 @@
 )
 
 
-class JsonRpcUser(HttpUser):
+class JrpcHttpUser(HttpUser):
     """Extension of HttpUser to provide JsonRPC support."""
 
     abstract = True

chainbench/user/protocol/ethereum.py

@@ -5,6 +5,7 @@
 
 from chainbench.test_data.ethereum import EthBeaconTestData
 from chainbench.user.http import HttpUser
+from chainbench.user.wss import WSSubscription
 
 logger = logging.getLogger(__name__)
 
@@ -329,3 +330,8 @@ class TestEthMethod(EthBeaconUser):
     @task
     def run_task(self) -> None:
         self.method_to_task_function(self.environment.parsed_options.method)()
+
+
+class EthSubscribe(WSSubscription):
+    def __init__(self, params: dict | list):
+        super().__init__("eth_subscribe", params, "eth_unsubscribe")

chainbench/user/protocol/evm.py

@@ -8,13 +8,13 @@
     Tx,
     TxHash,
 )
-from chainbench.user.jsonrpc import JsonRpcUser
+from chainbench.user.jsonrpc import JrpcHttpUser
 from chainbench.user.tag import tag
 from chainbench.util.jsonrpc import RpcCall
 from chainbench.util.rng import RNG, RNGManager
 
 
-class EvmBaseUser:
+class EvmBaseUser(JrpcHttpUser):
     abstract = True
     test_data: EvmTestData = EvmTestData()
     rng = RNGManager()
@@ -396,7 +396,7 @@ def web3_sha3(self) -> RpcCall:
         return RpcCall(method="web3_sha3", params=[self.test_data.get_random_tx_hash(self.rng.get_rng())])
 
 
-class EvmUser(EvmRpcMethods, JsonRpcUser):
+class EvmUser(EvmRpcMethods):
     abstract = True
 
     @staticmethod

chainbench/user/protocol/solana.py

@@ -4,12 +4,12 @@
 from solders.message import Message
 
 from chainbench.test_data import Account, BlockNumber, SolanaTestData, TxHash
-from chainbench.user.jsonrpc import JsonRpcUser
+from chainbench.user.jsonrpc import JrpcHttpUser
 from chainbench.util.jsonrpc import RpcCall
 from chainbench.util.rng import RNG, RNGManager
 
 
-class SolanaBaseUser(JsonRpcUser):
+class SolanaBaseUser(JrpcHttpUser):
     abstract = True
     test_data = SolanaTestData()
     rng = RNGManager()

chainbench/user/protocol/starknet.py

@@ -2,11 +2,11 @@
 
 from chainbench.test_data import StarkNetTestData
 from chainbench.test_data.blockchain import Account, TxHash
-from chainbench.user.jsonrpc import JsonRpcUser
+from chainbench.user.jsonrpc import JrpcHttpUser
 from chainbench.util.rng import RNG
 
 
-class StarkNetUser(JsonRpcUser):
+class StarkNetUser(JrpcHttpUser):
     abstract = True
     test_data = StarkNetTestData()
 

chainbench/user/wss.py

@@ -56,6 +56,9 @@ def __init__(self, environment: Environment):
         self._requests: dict[int, WSRequest] = {}
         self._running: bool = False
 
+    def get_subscription(self, subscription_id: str | int):
+        return self.subscriptions[self.subscription_ids_to_index[subscription_id]]
+
     @task
     def dummy_task(self):
         gevent.sleep(3600)