Update annotation in wrapper test files

Toni-SM · Aug 27, 2024 · 9fdb103 · 9fdb103
1 parent 48459c4
commit 9fdb103
Show file tree

Hide file tree

Showing 7 changed files with 9 additions and 8 deletions.
diff --git a/tests/jax/test_jax_wrapper_isaacgym.py b/tests/jax/test_jax_wrapper_isaacgym.py
@@ -58,7 +58,7 @@ def close(self) -> None:
 
 @pytest.mark.parametrize("backend", ["jax", "numpy"])
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, backend: str, num_states):
+def test_env(capsys: pytest.CaptureFixture, backend: str, num_states: str):
     config.jax.backend = backend
     Array = jax.Array if backend == "jax" else np.ndarray
 

diff --git a/tests/jax/test_jax_wrapper_isaaclab.py b/tests/jax/test_jax_wrapper_isaaclab.py
@@ -68,7 +68,7 @@ def close(self) -> None:
 
 @pytest.mark.parametrize("backend", ["jax", "numpy"])
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, backend: str, num_states):
+def test_env(capsys: pytest.CaptureFixture, backend: str, num_states: int):
     config.jax.backend = backend
     Array = jax.Array if backend == "jax" else np.ndarray
 

diff --git a/tests/jax/test_jax_wrapper_omniisaacgym.py b/tests/jax/test_jax_wrapper_omniisaacgym.py
@@ -52,7 +52,7 @@ def close(self) -> None:
 
 @pytest.mark.parametrize("backend", ["jax", "numpy"])
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, backend: str, num_states):
+def test_env(capsys: pytest.CaptureFixture, backend: str, num_states: int):
     config.jax.backend = backend
     Array = jax.Array if backend == "jax" else np.ndarray
 

diff --git a/tests/torch/test_torch_wrapper_gym.py b/tests/torch/test_torch_wrapper_gym.py
@@ -45,12 +45,13 @@ def test_env(capsys: pytest.CaptureFixture):
 
     env.close()
 
-def test_vectorized_env(capsys: pytest.CaptureFixture):
+@pytest.mark.parametrize("vectorization_mode", ["async", "sync"])
+def test_vectorized_env(capsys: pytest.CaptureFixture, vectorization_mode: str):
     num_envs = 10
     action = torch.ones((num_envs, 1))
 
     # load wrap the environment
-    original_env = gym.vector.make("Pendulum-v1", num_envs=num_envs, asynchronous=False)
+    original_env = gym.vector.make("Pendulum-v1", num_envs=num_envs, asynchronous=(vectorization_mode == "async"))
     env = wrap_env(original_env, "auto")
     assert isinstance(env, GymWrapper)
     env = wrap_env(original_env, "gym")

diff --git a/tests/torch/test_torch_wrapper_isaacgym.py b/tests/torch/test_torch_wrapper_isaacgym.py
@@ -54,7 +54,7 @@ def close(self) -> None:
 
 
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, num_states):
+def test_env(capsys: pytest.CaptureFixture, num_states: int):
     num_envs = 10
     action = torch.ones((num_envs, 1))
 

diff --git a/tests/torch/test_torch_wrapper_isaaclab.py b/tests/torch/test_torch_wrapper_isaaclab.py
@@ -64,7 +64,7 @@ def close(self) -> None:
 
 
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, num_states):
+def test_env(capsys: pytest.CaptureFixture, num_states: int):
     num_envs = 10
     action = torch.ones((num_envs, 1))
 

diff --git a/tests/torch/test_torch_wrapper_omniisaacgym.py b/tests/torch/test_torch_wrapper_omniisaacgym.py
@@ -48,7 +48,7 @@ def close(self) -> None:
 
 
 @pytest.mark.parametrize("num_states", [0, 5])
-def test_env(capsys: pytest.CaptureFixture, num_states):
+def test_env(capsys: pytest.CaptureFixture, num_states: int):
     num_envs = 10
     action = torch.ones((num_envs, 1))