Use Protocol instead of Callable in rate_limit_wrapper_test.

Unfortunately Python's type checking doesn't currently support functions with
attributes (python/mypy#2087) and it fails with a
message like `No attribute 'timestamp' on Callable[[Any], Any] [attribute-error]`.
This change adds a `Protocol` to the unit test such that the function can be
recognized as having the required attributes.

PiperOrigin-RevId: 477439369

Author: kenjitoyama

android_env/wrappers/rate_limit_wrapper_test.py

@@ -16,6 +16,7 @@
 """Tests for rate_limit_wrapper."""
 
 import time
+from typing import Any, List, Protocol
 from unittest import mock
 
 from absl.testing import absltest
@@ -46,6 +47,17 @@ def _get_base_env():
   return env
 
 
+class _FnWithTimestamps(Protocol):
+  """A function with `timestamp` and `timestamps` attributes."""
+
+  timestamp: float
+  timestamps: List[float]
+
+
+def _with_timestamp(fn: Any) -> _FnWithTimestamps:
+  return fn
+
+
 class RateLimitWrapperTest(parameterized.TestCase):
 
   @parameterized.named_parameters(
@@ -107,6 +119,7 @@ def test_enabled_sleep_type_before(self, mock_sleep):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
@@ -145,6 +158,7 @@ def test_enabled_sleep_type_after(self, mock_sleep):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
@@ -185,12 +199,14 @@ def test_enabled_sleep_type_after_with_repeat(self, mock_sleep):
     _ = wrapper.reset()
     mock_sleep.assert_not_called()  # It should never sleep during reset().
 
+    @_with_timestamp
     def _sleep_fn(sleep_time):
       _sleep_fn.timestamp = time.time()
       self.assertBetween(sleep_time, 0.0, 33.33)
 
     mock_sleep.side_effect = _sleep_fn
 
+    @_with_timestamp
     def _step_fn(action):
       # On even calls the action should be the actual agent action, but on odd
       # calls they should be REPEATs.