Fix time bugs in EngineProcessor.get_schedule (quantumlib#2878)

- Fix default "now" being assigned at init time instead of at call time
- Refactored method to take a datetime (absolute) or a timedelta (relative to now)
- Fixed specifying from_time but not to_time (or vice versa) resulting in a bad filter containing None
- Add 'freezegun' requirement for testing with control over 'now'

cirq/google/engine/engine_processor.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import datetime
 
-from typing import List, Optional, TYPE_CHECKING
+from typing import List, Optional, TYPE_CHECKING, Union, Tuple
 from pytz import utc
 
 from cirq.google.engine.client.quantum import types as qtypes
@@ -257,49 +257,110 @@ def update_reservation(self,
 
     def list_reservations(
             self,
-            from_time: datetime.datetime = datetime.datetime.now(tz=utc),
-            to_time: datetime.datetime = datetime.datetime.now(tz=utc) +
-            datetime.timedelta(weeks=2)) -> List[EngineTimeSlot]:
+            from_time: Union[None, datetime.datetime, datetime.
+                             timedelta] = datetime.timedelta(),
+            to_time: Union[None, datetime.datetime, datetime.
+                           timedelta] = datetime.timedelta(weeks=2)
+    ) -> List[EngineTimeSlot]:
         """Retrieves the reservations from a processor.
 
         Only reservations from this processor and project will be
-        returned.  The schedule may be filtered by starting and ending time.
+        returned. The schedule may be filtered by starting and ending time.
+
+        Args:
+            from_time: Filters the returned reservations to only include entries
+                that end no earlier than the given value. Specified either as an
+                absolute time (datetime.datetime) or as a time relative to now
+                (datetime.timedelta). Defaults to now (a relative time of 0).
+                Set to None to omit this filter.
+            to_time: Filters the returned reservations to only include entries
+                that start no later than the given value. Specified either as an
+                absolute time (datetime.datetime) or as a time relative to now
+                (datetime.timedelta). Defaults to two weeks from now (a relative
+                time of two weeks). Set to None to omit this filter.
+
+        Returns:
+            A list of reservations.
         """
-        filters = []
-        if from_time:
-            filters.append(f'start_time < {int(to_time.timestamp())}')
-        if to_time:
-            filters.append(f'end_time > {int(from_time.timestamp())}')
+        filters = _to_date_time_filters(from_time, to_time)
         filter_str = ' AND '.join(filters)
         return self.context.client.list_reservations(self.project_id,
                                                      self.processor_id,
                                                      filter_str)
 
     def get_schedule(
             self,
-            from_time: datetime.datetime = datetime.datetime.now(tz=utc),
-            to_time: datetime.datetime = datetime.datetime.now(tz=utc) +
-            datetime.timedelta(weeks=2),
-            time_slot_type: qenums.QuantumTimeSlot.TimeSlotType = None
+            from_time: Union[None, datetime.datetime, datetime.
+                             timedelta] = datetime.timedelta(),
+            to_time: Union[None, datetime.datetime, datetime.
+                           timedelta] = datetime.timedelta(weeks=2),
+            time_slot_type: Optional[qenums.QuantumTimeSlot.TimeSlotType] = None
     ) -> List[EngineTimeSlot]:
         """Retrieves the schedule for a processor.
 
         The schedule may be filtered by time.
 
         Time slot type will be supported in the future.
+
+        Args:
+            from_time: Filters the returned schedule to only include entries
+                that end no earlier than the given value. Specified either as an
+                absolute time (datetime.datetime) or as a time relative to now
+                (datetime.timedelta). Defaults to now (a relative time of 0).
+                Set to None to omit this filter.
+            to_time: Filters the returned schedule to only include entries
+                that start no later than the given value. Specified either as an
+                absolute time (datetime.datetime) or as a time relative to now
+                (datetime.timedelta). Defaults to two weeks from now (a relative
+                time of two weeks). Set to None to omit this filter.
+            time_slot_type: Filters the returned schedule to only include
+                entries with a given type (e.g. maintenance, open swim).
+                Defaults to None. Set to None to omit this filter.
+
+        Returns:
+            Schedule time slots.
         """
-        filters = []
-        if from_time:
-            filters.append(f'start_time < {int(to_time.timestamp())}')
-        if to_time:
-            filters.append(f'end_time > {int(from_time.timestamp())}')
-        if time_slot_type:
+        filters = _to_date_time_filters(from_time, to_time)
+        if time_slot_type is not None:
             filters.append(f'time_slot_type = {time_slot_type.name}')
         filter_str = ' AND '.join(filters)
         return self.context.client.list_time_slots(self.project_id,
                                                    self.processor_id,
                                                    filter_str)
 
-    def __str__(self) -> str:
-        return (f'EngineProcessor(project_id=\'{self.project_id}\', '
-                f'processor_id=\'{self.processor_id}\')')
+    def __str__(self):
+        return (f"EngineProcessor(project_id={self.project_id!r}, "
+                f"processor_id={self.processor_id!r})")
+
+
+def _to_date_time_filters(
+        from_time: Union[None, datetime.datetime, datetime.timedelta],
+        to_time: Union[None, datetime.datetime, datetime.timedelta]
+) -> List[str]:
+    now = datetime.datetime.now()
+
+    if from_time is None:
+        start_time = None
+    elif isinstance(from_time, datetime.timedelta):
+        start_time = now + from_time
+    elif isinstance(from_time, datetime.datetime):
+        start_time = from_time
+    else:
+        raise ValueError(
+            f"Don't understand from_time of type {type(from_time)}.")
+
+    if to_time is None:
+        end_time = None
+    elif isinstance(to_time, datetime.timedelta):
+        end_time = now + to_time
+    elif isinstance(to_time, datetime.datetime):
+        end_time = to_time
+    else:
+        raise ValueError(f"Don't understand to_time of type {type(to_time)}.")
+
+    filters = []
+    if end_time is not None:
+        filters.append(f'start_time < {int(end_time.timestamp())}')
+    if start_time is not None:
+        filters.append(f'end_time > {int(start_time.timestamp())}')
+    return filters

cirq/google/engine/engine_processor_test.py

@@ -14,7 +14,9 @@
 
 from unittest import mock
 import datetime
+
 import pytest
+import freezegun
 from google.protobuf.duration_pb2 import Duration
 from google.protobuf.text_format import Merge
 from google.protobuf.timestamp_pb2 import Timestamp
@@ -509,6 +511,94 @@ def test_get_schedule_filter_by_time_slot(list_time_slots):
         'time_slot_type = MAINTENANCE')
 
 
+@freezegun.freeze_time(datetime.datetime.utcfromtimestamp(100_000))
+@mock.patch('cirq.google.engine.engine_client.EngineClient.list_time_slots')
+def test_get_schedule_time_filter_behavior(list_time_slots):
+    list_time_slots.return_value = []
+    processor = cg.EngineProcessor('proj', 'p0', EngineContext())
+
+    processor.get_schedule()
+    list_time_slots.assert_called_with(
+        'proj', 'p0',
+        f'start_time < {100_000 + 60*60*24*14} AND end_time > {100_000}')
+
+    with pytest.raises(ValueError, match='from_time of type'):
+        processor.get_schedule(from_time=object())
+
+    with pytest.raises(ValueError, match='to_time of type'):
+        processor.get_schedule(to_time=object())
+
+    processor.get_schedule(from_time=None, to_time=None)
+    list_time_slots.assert_called_with('proj', 'p0', '')
+
+    processor.get_schedule(from_time=datetime.timedelta(0), to_time=None)
+    list_time_slots.assert_called_with('proj', 'p0', f'end_time > {100_000}')
+
+    processor.get_schedule(from_time=datetime.timedelta(seconds=200),
+                           to_time=None)
+    list_time_slots.assert_called_with('proj', 'p0', f'end_time > {100_200}')
+
+    processor.get_schedule(from_time=datetime.datetime.utcfromtimestamp(52),
+                           to_time=None)
+    list_time_slots.assert_called_with('proj', 'p0', f'end_time > {52}')
+
+    processor.get_schedule(from_time=None, to_time=datetime.timedelta(0))
+    list_time_slots.assert_called_with('proj', 'p0', f'start_time < {100_000}')
+
+    processor.get_schedule(from_time=None,
+                           to_time=datetime.timedelta(seconds=200))
+    list_time_slots.assert_called_with('proj', 'p0', f'start_time < {100_200}')
+
+    processor.get_schedule(from_time=None,
+                           to_time=datetime.datetime.utcfromtimestamp(52))
+    list_time_slots.assert_called_with('proj', 'p0', f'start_time < {52}')
+
+
+@freezegun.freeze_time(datetime.datetime.utcfromtimestamp(100_000))
+@mock.patch('cirq.google.engine.engine_client.EngineClient.list_reservations')
+def test_list_reservations_time_filter_behavior(list_reservations):
+    list_reservations.return_value = []
+    processor = cg.EngineProcessor('proj', 'p0', EngineContext())
+
+    processor.list_reservations()
+    list_reservations.assert_called_with(
+        'proj', 'p0',
+        f'start_time < {100_000 + 60*60*24*14} AND end_time > {100_000}')
+
+    with pytest.raises(ValueError, match='from_time of type'):
+        processor.list_reservations(from_time=object())
+
+    with pytest.raises(ValueError, match='to_time of type'):
+        processor.list_reservations(to_time=object())
+
+    processor.list_reservations(from_time=None, to_time=None)
+    list_reservations.assert_called_with('proj', 'p0', '')
+
+    processor.list_reservations(from_time=datetime.timedelta(0), to_time=None)
+    list_reservations.assert_called_with('proj', 'p0', f'end_time > {100_000}')
+
+    processor.list_reservations(from_time=datetime.timedelta(seconds=200),
+                                to_time=None)
+    list_reservations.assert_called_with('proj', 'p0', f'end_time > {100_200}')
+
+    processor.list_reservations(
+        from_time=datetime.datetime.utcfromtimestamp(52), to_time=None)
+    list_reservations.assert_called_with('proj', 'p0', f'end_time > {52}')
+
+    processor.list_reservations(from_time=None, to_time=datetime.timedelta(0))
+    list_reservations.assert_called_with('proj', 'p0',
+                                         f'start_time < {100_000}')
+
+    processor.list_reservations(from_time=None,
+                                to_time=datetime.timedelta(seconds=200))
+    list_reservations.assert_called_with('proj', 'p0',
+                                         f'start_time < {100_200}')
+
+    processor.list_reservations(from_time=None,
+                                to_time=datetime.datetime.utcfromtimestamp(52))
+    list_reservations.assert_called_with('proj', 'p0', f'start_time < {52}')
+
+
 def test_str():
     processor = cg.EngineProcessor('a', 'p', EngineContext())
     assert str(

dev_tools/conf/mypy.ini

@@ -5,7 +5,7 @@ follow_imports = silent
 ignore_missing_imports = true
 
 # 3rd-party libs for which we don't have stubs
-[mypy-apiclient.*,matplotlib.*,mpl_toolkits,multiprocessing.dummy,numpy.*,oauth2client.*,pandas.*,pytest.*,scipy.*,sortedcontainers.*,setuptools.*,pylatex.*,networkx.*,qiskit.*,pypandoc.*,ply.*,_pytest.*,google.api.*,google.api_core.*,grpc.*,google.oauth2.*,google.protobuf.text_format.*]
+[mypy-apiclient.*,freezegun.*,matplotlib.*,mpl_toolkits,multiprocessing.dummy,numpy.*,oauth2client.*,pandas.*,pytest.*,scipy.*,sortedcontainers.*,setuptools.*,pylatex.*,networkx.*,qiskit.*,pypandoc.*,ply.*,_pytest.*,google.api.*,google.api_core.*,grpc.*,google.oauth2.*,google.protobuf.text_format.*]
 follow_imports = silent
 ignore_missing_imports = true
 

requirements.txt

@@ -1,6 +1,7 @@
 # Runtime requirements for the python 3 version of cirq.
 
 dataclasses; python_version < '3.7'
+freezegun~=0.3.15
 google-api-core[grpc] >= 1.14.0, < 2.0.0dev
 matplotlib~=3.0
 networkx~=2.4