Allow deprecation for freezegun (#3950)

Allows usage of deprecated functionality for test functions leveraging freezegun.

During cirq.google migration this came up - as we don't have yet any deprecated attributes it didn't bite us. 

Note: all other alternatives are significantly more complex, for example
- we could make this more precise and introduce an env var specifically for the attributes
- we could create a wrapper that removes all deprecated attributes before running these tests

cirq/google/engine/engine_processor_test.py

@@ -554,6 +554,36 @@ def test_get_schedule_filter_by_time_slot(list_time_slots):
     )
 
 
+def _allow_deprecated_freezegun(func):
+    # a local hack, as freeze_time walks through all the sys.modules, and retrieves all the
+    # attributes for all modules when it reaches deprecated module attributes, we throw an error
+    # as the deprecation module thinks Cirq is using something deprecated. This hack SHOULD NOT be
+    # used elsewhere, it is specific to freezegun functionality.
+    def wrapper(*args, **kwargs):
+        import os
+        from cirq.testing.deprecation import ALLOW_DEPRECATION_IN_TEST
+
+        orig_exist, orig_value = (
+            ALLOW_DEPRECATION_IN_TEST in os.environ,
+            os.environ.get(ALLOW_DEPRECATION_IN_TEST, None),
+        )
+
+        os.environ[ALLOW_DEPRECATION_IN_TEST] = 'True'
+        try:
+            return func(*args, **kwargs)
+        finally:
+            if orig_exist:
+                # mypy can't resolve that orig_exist ensures that orig_value
+                # of type Optional[str] can't be None
+                # coverage: ignore
+                os.environ[ALLOW_DEPRECATION_IN_TEST] = orig_value  # type: ignore
+            else:
+                del os.environ[ALLOW_DEPRECATION_IN_TEST]
+
+    return wrapper
+
+
+@_allow_deprecated_freezegun
 @freezegun.freeze_time()
 @mock.patch('cirq.google.engine.engine_client.EngineClient.list_time_slots')
 def test_get_schedule_time_filter_behavior(list_time_slots):
@@ -597,6 +627,7 @@ def test_get_schedule_time_filter_behavior(list_time_slots):
     list_time_slots.assert_called_with('proj', 'p0', f'start_time < {utc_ts}')
 
 
+@_allow_deprecated_freezegun
 @freezegun.freeze_time()
 @mock.patch('cirq.google.engine.engine_client.EngineClient.list_reservations')
 def test_list_reservations_time_filter_behavior(list_reservations):

cirq/testing/deprecation.py

@@ -51,6 +51,10 @@ def assert_deprecated(*msgs: str, deadline: str, count: Optional[int] = 1):
             messages have to equal count.
     """
 
+    orig_exist, orig_value = (
+        ALLOW_DEPRECATION_IN_TEST in os.environ,
+        os.environ.get(ALLOW_DEPRECATION_IN_TEST, None),
+    )
     os.environ[ALLOW_DEPRECATION_IN_TEST] = 'True'
     try:
         with assert_logs(
@@ -62,7 +66,12 @@ def assert_deprecated(*msgs: str, deadline: str, count: Optional[int] = 1):
             yield True
     finally:
         try:
-            del os.environ[ALLOW_DEPRECATION_IN_TEST]
+            if orig_exist:
+                # mypy can't resolve that orig_exist ensures that orig_value
+                # of type Optional[str] can't be None
+                os.environ[ALLOW_DEPRECATION_IN_TEST] = orig_value  # type: ignore
+            else:
+                del os.environ[ALLOW_DEPRECATION_IN_TEST]
         except:
             # this is only for nested deprecation checks
             pass