Refactor, prettify, lints, docstings.

functional-tests/entry.py

@@ -16,7 +16,8 @@
 from factory import factory
 from utils import *
 from utils.constants import *
-
+from load.cfg import RethLoadConfigBuilder
+from load.jobs import EthJob
 
 TEST_DIR: str = "tests"
 
@@ -87,7 +88,7 @@ def main(argv):
     reth_fac = factory.RethFactory([12600 + i for i in range(100 * 3)])
     prover_client_fac = factory.ProverClientFactory([12900 + i for i in range(100 * 3)])
     bridge_client_fac = factory.BridgeClientFactory([13200 + i for i in range(100)])
-    load_gen_fac = factory.LoadGeneratorFactory([14000 + i for i in range(100)])
+    load_gen_fac = factory.LoadGeneratorFactory([13300 + i for i in range(100)])
     seq_signer_fac = factory.StrataSequencerFactory()
 
     factories = {
@@ -101,6 +102,9 @@ def main(argv):
         "load_generator": load_gen_fac,
     }
 
+    reth_load_env = testenv.LoadEnvConfig()
+    reth_load_env.with_load_builder(RethLoadConfigBuilder().with_jobs([EthJob]).with_rate(15))
+
     global_envs = {
         # Basic env is the default env for all tests.
         "basic": testenv.BasicEnvConfig(101),
@@ -116,7 +120,7 @@ def main(argv):
             2
         ),  # TODO: Need to generate at least horizon blocks, based on params
         "prover": testenv.BasicEnvConfig(101),
-        "load": testenv.LoadEnvConfig(),
+        "load_reth": reth_load_env,
     }
 
     setup_root_logger()

functional-tests/envs/testenv.py

@@ -10,10 +10,8 @@
 
 from envs.rollup_params_cfg import RollupConfig
 from load.cfg import LoadConfig, LoadConfigBuilder
-from load.jobs.reth import EthJob
 from utils import *
 from utils.constants import *
-from typing import Callable
 
 
 class StrataTester(flexitest.Test):
@@ -270,7 +268,6 @@ def init(self, ctx: flexitest.EnvContext) -> flexitest.LiveEnv:
         return BasicLiveEnv(svcs, bridge_pk, rollup_cfg)
 
 
-# TODO: refactor out the common code from all envs as this is becoming messy.
 class HubNetworkEnvConfig(flexitest.EnvConfig):
     def __init__(
         self,
@@ -409,29 +406,27 @@ def init(self, ctx: flexitest.EnvContext) -> flexitest.LiveEnv:
         return BasicLiveEnv(svcs, bridge_pk, rollup_cfg)
 
 
-# TODO: refactor out the common code from all envs as this is becoming messy.
-# Currently, it's not ideal either, as it heavily depends on BasicEnvConfig.
-# TODO(load): make it more generic - right now it targets reth service and has hardcoded job.
+# TODO: Maybe, we need to make it dynamic to enhance any EnvConfig with load testing capabilities.
 class LoadEnvConfig(BasicEnvConfig):
-    _load_config_builder: Callable[[dict[str, flexitest.Service]], LoadConfig]
+    _load_cfgs: list[Callable[[dict[str, flexitest.Service]], LoadConfig]] = []
 
-    def __init__(
-        self, builder: Callable[[dict[str, flexitest.Service]], LoadConfig], *args, **kwargs
-    ):
-        super().__init__(*args, **kwargs)
-        self._load_config_builder = builder
+    def with_load_builder(self, builder: LoadConfigBuilder):
+        self._load_cfgs.append(builder)
+        return self
 
     def init(self, ctx: flexitest.EnvContext) -> flexitest.LiveEnv:
         basic_live_env = super().init(ctx)
 
-        svcs = basic_live_env.svcs
-        reth = svcs["reth"]
-        web3_port = reth.get_prop("eth_rpc_http_port")
-
-        load_cfg = self._load_config_builder(svcs)
+        if not self._load_cfgs:
+            raise Exception(
+                "LoadEnv has no load builders! Specify load builders or just use BasicEnv."
+            )
 
-        load_cfg: LoadConfig = LoadConfig([EthJob], f"http://localhost:{web3_port}", 50)
+        # Create load generator services for all the builders.
+        svcs = basic_live_env.svcs
         load_fac = ctx.get_factory("load_generator")
-        svcs["load_generator"] = load_fac.create_simple_loadgen(load_cfg)
+        for builder in self._load_cfgs:
+            load_cfg: LoadConfig = builder(svcs)
+            svcs[f"load_generator.{builder.name}"] = load_fac.create_simple_loadgen(load_cfg)
 
         return BasicLiveEnv(svcs, basic_live_env._bridge_pk, basic_live_env._rollup_cfg)

functional-tests/load/cfg.py

@@ -1,31 +1,76 @@
-from locust import HttpUser
+from dataclasses import dataclass
 
+import flexitest
 
+from load.job import StrataLoadJob
+
+
+@dataclass(frozen=True)
 class LoadConfig:
     """
-    A data class holding host and job classes that describe load requests.
+    Config for the load service.
     """
 
-    user_classes: list[HttpUser]
-    host: str
-    spawn_rate: int
+    jobs: list[StrataLoadJob]
+    """A set of jobs that emit the load towards the host."""
 
-    def __init__(self, user_classes: list[HttpUser], host: str, spawn_rate: int):
-        # "Patch" user_classes with a given host
-        for user_cls in user_classes:
-            user_cls.host = host
+    host: str
+    """The host that will accept load requests."""
 
-        self.user_classes = user_classes
-        self.host = host
-        self.spawn_rate = spawn_rate
+    spawn_rate: int
+    """The rate at which all the jobs will emit the load."""
 
 
 class LoadConfigBuilder:
     """
-    A builder for the load config.
+    An abstract builder of the `LoadConfig`.
     """
 
+    jobs: list[StrataLoadJob] = []
+    """A set of jobs that emit the load towards the host."""
+
+    spawn_rate: int = 10
+    """The rate at which all the jobs will emit the load."""
+
+    service_name: str | None = None
+    """The name of the service to emit the load."""
+
     def __init__(self):
-        pass
+        if not self.service_name:
+            raise Exception("LoadConfigBuilder: missing service_name attribute.")
+
+    def with_jobs(self, jobs: list[StrataLoadJob]):
+        self.jobs.extend(jobs)
+        return self
+
+    def with_rate(self, rate: int):
+        self.spawn_rate = rate
+        return self
+
+    def __call__(self, svcs) -> LoadConfig:
+        if not self.jobs:
+            raise Exception("LoadConfigBuilder: load jobs list is empty")
+
+        host = self.host_url(svcs)
+        # Patch jobs by the host.
+        for job in self.jobs:
+            job.host = host
+
+        return LoadConfig(self.jobs, host, self.spawn_rate)
+
+    def host_url(self, _svcs: dict[str, flexitest.Service]) -> str:
+        raise NotImplementedError()
+
+    @property
+    def name(self):
+        return self.service_name
+
+
+class RethLoadConfigBuilder(LoadConfigBuilder):
+    service_name: str = "reth"
+    spawn_rate: int = 20
 
-    def __call__(self, )
+    def host_url(self, svcs: dict[str, flexitest.Service]) -> str:
+        reth = svcs["reth"]
+        web3_port = reth.get_prop("eth_rpc_http_port")
+        return f"http://localhost:{web3_port}"

functional-tests/load/job.py

@@ -2,11 +2,19 @@
 import web3.middleware
 from locust import HttpUser
 
-FUND_AMOUNT = 1_000_000_000_000_000_000_000  # 1000 ETH
+
+class StrataLoadJob(HttpUser):
+    """
+    A common layer for all the load jobs in the load tests.
+    """
+
+    pass
 
 
 # TODO(load): configure the structured logging as we do in the tests.
-class BaseEthJob(HttpUser):
+class BaseRethLoadJob(StrataLoadJob):
+    fund_amount: int = 1_000_000_000_000_000_000_000  # 1000 ETH
+
     def on_start(self):
         root_w3, genesis_acc = self.w3_with_genesis_acc()
         self._root_w3 = root_w3
@@ -46,9 +54,8 @@ def _init_w3(self, init):
     def _fund_account(self, acc):
         print(f"FUNDING ACCOUNT {acc}")
         source = self._root_w3.address
-        to_transfer = FUND_AMOUNT
         tx_hash = self._root_w3.eth.send_transaction(
-            {"to": acc, "value": hex(to_transfer), "gas": hex(100000), "from": source}
+            {"to": acc, "value": hex(self.fund_amount), "gas": hex(100000), "from": source}
         )
 
         tx_receipt = self._root_w3.eth.wait_for_transaction_receipt(tx_hash, timeout=120)

functional-tests/load/jobs/__init__.py

@@ -0,0 +1 @@
+from .reth import *

functional-tests/load/jobs/reth.py

@@ -1,9 +1,9 @@
 from locust import task
 
-from load.job import BaseEthJob
+from load.job import BaseRethLoadJob
 
 
-class EthJob(BaseEthJob):
+class EthJob(BaseRethLoadJob):
     def on_start(self):
         super().on_start()
 

functional-tests/load/service.py

@@ -13,21 +13,24 @@
 # TODO(load): enhance it to be able to increase/decrease the load from test runtime.
 class LoadGeneratorService(flexitest.Service):
     """
-    A separate flexitest service that is able to generate load as specified in the LoadConfig.
+    A separate flexitest service that is able to generate the load as specified by `LoadConfig`.
     """
 
-    # Locust Environment for running the load.
     env: Environment
-    # Locust local runner that actually dispatches the load.
+    """Locust Environment for running the load against."""
+
     runner: LocalRunner
-    # A thin config that specifies load params: jobs, host, rate, etc.
+    """Locust local runner that actually dispatches the load."""
+
     cfg: LoadConfig
-    # Whether the LoadGenerator (as a flexitest service) is started.
+    """A config that specifies load params: jobs, host, rate, etc."""
+
     _is_started: bool
+    """Whether the LoadGenerator (as a flexitest service) is started."""
 
     def __init__(self, logfile, cfg: LoadConfig):
         self._is_started = False
-        self.env = Environment(user_classes=cfg.user_classes, events=events)
+        self.env = Environment(user_classes=cfg.jobs, events=events)
         self.runner = self.env.create_local_runner()
         self.cfg = cfg
 

functional-tests/tests/load/basic_load.py

@@ -10,11 +10,11 @@
 @flexitest.register
 class BasicLoadGenerationTest(testenv.StrataTester):
     def __init__(self, ctx: flexitest.InitContext):
-        ctx.set_env("load")
+        ctx.set_env("load_reth")
 
     def main(self, ctx: flexitest.RunContext):
         print("test is running")
-        load = ctx.get_service("load_generator")
+        load = ctx.get_service("load_generator.reth")
         _loadrpc = load.create_rpc()
         time.sleep(5)