diff --git a/tests/e2e_integration_tests/action_masking/test_agents_use_action_masks.py b/tests/e2e_integration_tests/action_masking/test_agents_use_action_masks.py
index 4260c605..addf6dca 100644
--- a/tests/e2e_integration_tests/action_masking/test_agents_use_action_masks.py
+++ b/tests/e2e_integration_tests/action_masking/test_agents_use_action_masks.py
@@ -1,6 +1,7 @@
 # © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
 from typing import Dict
 
+import pytest
 import yaml
 from ray.rllib.algorithms.ppo import PPOConfig
 from ray.rllib.core.rl_module.marl_module import MultiAgentRLModuleSpec