From 14af14269244bbd5f5734a80d41df53a17a21df8 Mon Sep 17 00:00:00 2001
From: Alessandro Palmas <alessandropalmas.mail@gmail.com>
Date: Sun, 10 Mar 2024 10:10:44 -0400
Subject: [PATCH] Add wrappers speed tests

---
 tests/test_speed.py | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/tests/test_speed.py b/tests/test_speed.py
index 60245fb..8544bdb 100644
--- a/tests/test_speed.py
+++ b/tests/test_speed.py
@@ -96,4 +96,31 @@ def test_speed_wrappers_mock(n_players, mocker):
 @pytest.mark.parametrize("n_players", n_players)
 def test_speed_gym_integration(game_id, n_players, mocker):
     use_mocker = False
-    assert func(game_id, n_players, WrappersSettings(), use_mocker, mocker) == 0
\ No newline at end of file
+    assert func(game_id, n_players, WrappersSettings(), use_mocker, mocker) == 0
+
+@pytest.mark.parametrize("game_id", game_ids)
+@pytest.mark.parametrize("n_players", n_players)
+def test_speed_wrappers_integration(game_id, n_players, mocker):
+    use_mocker = False
+
+    # Env wrappers settings
+    wrappers_settings = WrappersSettings()
+    wrappers_settings.no_op_max = 0
+    wrappers_settings.repeat_action = 1
+    wrappers_settings.normalize_reward = True
+    wrappers_settings.clip_reward = False
+    wrappers_settings.stack_frames = 4
+    wrappers_settings.dilation = 1
+    wrappers_settings.add_last_action = True
+    wrappers_settings.stack_actions = 12
+    wrappers_settings.scale = True
+    wrappers_settings.role_relative = True
+    wrappers_settings.flatten = True
+
+    suffix = ""
+    if n_players == 2:
+        suffix = "agent_0_"
+    wrappers_settings.filter_keys = ["stage", "timer", suffix + "own_side", suffix + "opp_side",
+                                     suffix + "opp_character", suffix + "action"]
+
+    assert func(game_id, n_players, wrappers_settings, use_mocker, mocker) == 0
\ No newline at end of file