fix bug in beacon user

chainstacklabs · Oct 28, 2024 · bdb0120 · bdb0120
1 parent e8d0296
commit bdb0120
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/chainbench/main.py b/chainbench/main.py
@@ -252,12 +252,12 @@ def start(
         sys.exit(1)
 
     if test_by_directory:
-        from locust.argument_parser import find_locustfiles
+        from locust.argument_parser import parse_locustfile_paths
         from locust.util.load_locustfile import load_locustfile
 
         user_classes = {}
         test_data_types = set()
-        for locustfile in find_locustfiles([profile_path.__str__()], True):
+        for locustfile in parse_locustfile_paths([profile_path.__str__()]):
             _, _user_classes, _ = load_locustfile(locustfile)
             for key, value in _user_classes.items():
                 user_classes[key] = value

diff --git a/chainbench/user/protocol/ethereum.py b/chainbench/user/protocol/ethereum.py
@@ -110,7 +110,10 @@ def task_to_method(task_name: str) -> str:
 
     @classmethod
     def method_to_task_function(cls, method: str) -> t.Callable:
-        return getattr(cls, f"{method}_task")
+        if not method.endswith("task"):
+            return getattr(cls, f"{method}_task")
+        else:
+            return getattr(cls, method)
 
     def eth_v1_beacon_states_head_fork_task(self):
         self.get(