linting and formatting

Author: erwin-wee

chainbench/profile/solana/get_block.py

@@ -98,8 +98,10 @@ def make_rpc_call(
         path: str = "",
     ) -> None:
         """Make a JSON-RPC call."""
-        if rpc_call is None:
+        if rpc_call is None and method is not None:
             rpc_call = RpcCall(method, params)
+        else:
+            raise ValueError("Either rpc_call or method must be provided")
 
         if name == "" and method is not None:
             name = method

chainbench/user/jsonrpc.py

@@ -16,7 +16,7 @@
 
 class EvmBaseUser:
     abstract = True
-    test_data = EvmTestData()
+    test_data: EvmTestData = EvmTestData()
     rng = RNGManager()
 
     _default_trace_timeout = "120s"

chainbench/user/protocol/evm.py

@@ -5,18 +5,15 @@
 import orjson as json
 from gevent import Greenlet, Timeout
 from locust import User, task
+from locust.env import Environment
 from orjson import JSONDecodeError
 from websocket import WebSocket, WebSocketConnectionClosedException, create_connection
+
 from chainbench.util.jsonrpc import RpcCall
 
 
 class WSSubscription:
-    def __init__(
-            self,
-            subscribe_method: str,
-            subscribe_params: dict | list,
-            unsubscribe_method: str
-    ):
+    def __init__(self, subscribe_method: str, subscribe_params: dict | list, unsubscribe_method: str):
         self.subscribe_rpc_call: RpcCall = RpcCall(subscribe_method, subscribe_params)
         self.unsubscribe_method: str = unsubscribe_method
         self.subscribed: bool = False
@@ -38,7 +35,7 @@ def subscription_id(self):
 
 
 class WSRequest:
-    def __init__(self, rpc_call: RpcCall, start_time: int, subscription_index: int = None):
+    def __init__(self, rpc_call: RpcCall, start_time: int, subscription_index: int | None = None):
         self.rpc_call = rpc_call
         self.start_time = start_time
         self.subscription_index = subscription_index
@@ -52,7 +49,7 @@ class WssJrpcUser(User):
     subscriptions: list[WSSubscription] = []
     subscription_ids_to_index: dict[str | int, int] = {}
 
-    def __init__(self, environment):
+    def __init__(self, environment: Environment):
         super().__init__(environment)
         self._ws: WebSocket | None = None
         self._ws_greenlet: Greenlet | None = None
@@ -108,7 +105,7 @@ def get_notification_name(self, parsed_response: dict):
         # Override this method to return the name of the notification if this is not correct
         return parsed_response["method"]
 
-    def on_message(self, message):
+    def on_message(self, message: str | bytes):
         try:
             parsed_json: dict = json.loads(message)
             if "error" in parsed_json:
@@ -181,15 +178,24 @@ def receive_loop(self):
             self.logger.error("Connection closed by server, trying to reconnect...")
             self.on_start()
 
-    def send(self, rpc_call: RpcCall = None, method: str = None, params: dict | list = None, subscription_index: int = None):
+    def send(
+        self,
+        rpc_call: RpcCall | None = None,
+        method: str | None = None,
+        params: dict | list | None = None,
+        subscription_index: int | None = None,
+    ):
+        def _get_args():
+            if rpc_call:
+                return rpc_call
+            elif method:
+                return RpcCall(method, params)
+            else:
+                raise ValueError("Either rpc_call or method must be provided")
+
+        rpc_call = _get_args()
         self.logger.debug(f"Sending: {rpc_call or method}")
-        rpc = {
-            (None, None): None,
-            (None, method): RpcCall(method, params),
-            (rpc_call, None): rpc_call,
-        }
 
-        rpc_call = rpc[(rpc_call, method)]
         if rpc_call is None:
             raise ValueError("Either rpc_call or method must be provided")
 
@@ -202,4 +208,5 @@ def send(self, rpc_call: RpcCall = None, method: str = None, params: dict | list
         )
         json_body = json.dumps(rpc_call.request_body())
         self.logger.debug(f"WSReq: {json_body}")
-        self._ws.send(json_body)
+        if self._ws:
+            self._ws.send(json_body)

chainbench/user/wss.py

@@ -1,10 +1,11 @@
 import random
 import typing as t
+
 import orjson as json
 
 
 class RpcCall:
-    def __init__(self, method: str, params: list[t.Any] | dict | None = None, request_id: int = None) -> None:
+    def __init__(self, method: str, params: list[t.Any] | dict | None = None, request_id: int | None = None) -> None:
         self._request_id = request_id
         self.method = method
         self.params = params
@@ -15,7 +16,7 @@ def request_id(self) -> int:
             self._request_id = random.Random().randint(1, 100000000)
         return self._request_id
 
-    def request_body(self, request_id: int = None) -> dict:
+    def request_body(self, request_id: int | None = None) -> dict:
         """Generate a JSON-RPC request body."""
         if self.params is None:
             self.params = []
@@ -36,12 +37,7 @@ def request_body(self, request_id: int = None) -> dict:
 
 def generate_batch_request_body(rpc_calls: list[RpcCall]) -> str:
     """Generate a batch JSON-RPC request body."""
-    return json.dumps(
-        [
-            rpc_calls[i].request_body(i)
-            for i in range(1, len(rpc_calls))
-        ]
-    ).decode("utf-8")
+    return json.dumps([rpc_calls[i].request_body(i) for i in range(1, len(rpc_calls))]).decode("utf-8")
 
 
 def expand_rpc_calls(rpc_calls_weighted: dict[t.Callable[[], RpcCall], int]) -> list[RpcCall]: