feat: update tests

diambra · Apr 23, 2024 · 099920a · 099920a
1 parent 890bd69
commit 099920a
Showing 1 changed file with 19 additions and 13 deletions.
diff --git a/tests/test_sheeprl.py b/tests/test_sheeprl.py
@@ -30,7 +30,7 @@ def _test_agent(mocker, agent, kwargs):
 
     agent = importlib.import_module(f"agent-{agent}")
     os.environ["DIAMBRA_ENVS"] = "127.0.0.1:32781"
-    agent.main(**kwargs)
+    return agent.main(**kwargs)
 
 
 def _test_train_eval(
@@ -154,30 +154,36 @@ def test_sheeprl_evaluation(mocker):
 
 def test_sheeprl_ppo_agent(mocker):
     cfg_path = os.path.join(
-        ROOT_DIR, "/fake-logs/runs/ppo/doapp/fake-experiment/version_0/config.yaml"
+        ROOT_DIR, "example-logs/runs/ppo/doapp/experiment/version_0/config.yaml"
     )
     checkpoint_path = os.path.join(
         ROOT_DIR,
-        "/fake-logs/runs/ppo/doapp/fake-experiment/version_0/checkpoint/ckpt_1024_0.ckpt",
+        "example-logs/runs/ppo/doapp/experiment/version_0/checkpoint/ckpt_1024_0.ckpt",
     )
-    assert _test_agent(
-        mocker,
-        "ppo",
-        {"cfg_path": cfg_path, "checkpoint_path": checkpoint_path, "test": True},
+    assert (
+        _test_agent(
+            mocker,
+            "ppo",
+            {"cfg_path": cfg_path, "checkpoint_path": checkpoint_path, "test": True},
+        )
+        == 0
     )
 
 
 def test_sheeprl_dreamer_v3_agent(mocker):
     cfg_path = os.path.join(
         ROOT_DIR,
-        "/fake-logs/runs/dreamer_v3/doapp/fake-experiment/version_0/config.yaml",
+        "example-logs/runs/dreamer_v3/doapp/experiment/version_0/config.yaml",
     )
     checkpoint_path = os.path.join(
         ROOT_DIR,
-        "/fake-logs/runs/dreamer_v3/doapp/fake-experiment/version_0/checkpoint/ckpt_1024_0.ckpt",
+        "example-logs/runs/dreamer_v3/doapp/experiment/version_0/checkpoint/ckpt_1024_0.ckpt",
     )
-    assert _test_agent(
-        mocker,
-        "dreamer_v3",
-        {"cfg_path": cfg_path, "checkpoint_path": checkpoint_path, "test": True},
+    assert (
+        _test_agent(
+            mocker,
+            "dreamer_v3",
+            {"cfg_path": cfg_path, "checkpoint_path": checkpoint_path, "test": True},
+        )
+        == 0
     )