code: def test_resets_after_limit(self):
      max_steps = 5
      base_env = mock.MagicMock()
      wrapped_env = atari_wrappers.AtariTimeLimit(base_env, max_steps)
  
      base_env.gym.game_over = False
      base_env.reset.return_value = ts.restart(
          1
      )  # pytype: disable=wrong-arg-types
      base_env.step.return_value = ts.transition(
          2, 0
      )  # pytype: disable=wrong-arg-types
      action = 1
  
      for _ in range(max_steps + 1):
          wrapped_env.step(action)  # pytype: disable=wrong-arg-types
  
      self.assertTrue(wrapped_env.game_over)
      self.assertEqual(1, base_env.reset.call_count)
  
      wrapped_env.step(action)  # pytype: disable=wrong-arg-types
      self.assertFalse(wrapped_env.game_over)
      self.assertEqual(2, base_env.reset.call_count)
  
file_path: agents/tf_agents/environments/atari_wrappers_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing.absltest import mock
  from tf_agents.environments import atari_wrappers
  from tf_agents.trajectories import time_step
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the AtariTimeLimitTest obj's underlying member method code with example usage: obj.test_resets_after_limit() without return types