Fix tests

Author: MarcCote

tests/agents/test_agents.py

@@ -285,8 +285,8 @@ def test_create_agent():
     class TestRegisteredAgent(BaseAgent):
         name = "test_registered"
 
-        def __init__(self, args, env, **kwargs):
-            super().__init__(args, env, **kwargs)
+        def __init__(self, agent_args, env, **kwargs):
+            super().__init__(agent_args, env, **kwargs)
 
     # Clear and setup registry
     original_registry = AGENT_REGISTRY.copy()

tests/gym/envs/test_r2egym.py

@@ -23,6 +23,7 @@ def test_load_dataset(get_r2egym_env):
     assert sorted(task_data.keys()) == sorted(
         [
             "commit_hash",
+            "env_type",
             "docker_image",
             "execution_result_content",
             "expected_output_json",
@@ -75,6 +76,7 @@ def test_load_dataset_from_parquet(tmp_path):
     assert sorted(dataset_entry) == sorted(
         [
             "commit_hash",
+            "env_type",
             "docker_image",
             "execution_result_content",
             "expected_output_json",

tests/gym/envs/test_swe_bench.py

@@ -108,6 +108,7 @@ def test_load_dataset(get_swe_bench_env):
     assert sorted(task_data.keys()) == sorted(
         [
             "repo",
+            "env_type",
             "instance_id",
             "base_commit",
             "patch",

tests/gym/envs/test_swe_smith.py

@@ -25,6 +25,7 @@ def test_load_dataset(get_swe_smith_env):
     assert sorted(task_data.keys()) == sorted(
         [
             "instance_id",
+            "env_type",
             "repo",
             "patch",
             "FAIL_TO_PASS",
@@ -65,6 +66,7 @@ def test_load_dataset_from_parquet(tmp_path):
     assert sorted(dataset_entry.keys()) == sorted(
         [
             "instance_id",
+            "env_type",
             "repo",
             "patch",
             "FAIL_TO_PASS",
@@ -248,6 +250,7 @@ def test_running_solution_agent(get_swe_smith_env, tmp_path):
         "memory_size": 8,
         "max_steps": 1,
         "max_rewrite_steps": 1,
+        "env": env,
     }
     for tool_name in ["pdb", "eval", "submit"]:
         env.add_tool(Toolbox.get_tool(tool_name))