From 32b7e84b0ddf9f57ec1ff73fadc488543111166b Mon Sep 17 00:00:00 2001
From: Lingyi Hu <lingyihuu@gmail.com>
Date: Tue, 9 Feb 2021 19:47:42 +0800
Subject: [PATCH] address comments

---
 recirq/engine_utils.py                        | 16 ++++++++------
 recirq/engine_utils_test.py                   | 22 +++++++++++++++++--
 .../quantum_chess/experiments/batch_moves.py  |  3 +--
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/recirq/engine_utils.py b/recirq/engine_utils.py
index 822cc005..e3a65805 100644
--- a/recirq/engine_utils.py
+++ b/recirq/engine_utils.py
@@ -382,12 +382,18 @@ def _get_current_time():
 
 
 def get_available_processors(processor_names: List[str]):
-    """Checks the reservation status of the processors and returns a list of
+    """Returns a list of available processors.
+
+    Checks the reservation status of the processors and returns a list of
     processors that are available to run on at the present time.
+
+    Args:
+        processor_names: A list of processor names which are keys from QUANTUM_PROCESSORS.
     """
     project_id = os.environ['GOOGLE_CLOUD_PROJECT']
     engine = cirq.google.get_engine()
     available_processors = []
+    current_time = _get_current_time()
     for processor_name in processor_names:
         processor_id = get_processor_id_by_device_name(processor_name)
         if processor_id is None:
@@ -401,13 +407,9 @@ def get_available_processors(processor_names: List[str]):
             except ValueError:
                 continue
             # Ignore time slots that do not contain the current time.
-            if time_slot.start_time < _get_current_time() < time_slot.end_time:
+            if time_slot.start_time < current_time < time_slot.end_time:
                 # Time slots need to be either in OPEN_SWIM or reserved by the
                 # current project to be considered available.
-                if time_slot.slot_type == enums.QuantumTimeSlot.TimeSlotType.OPEN_SWIM:
-                    available_processors.append(processor_name)
-                    break
-                if time_slot.slot_type == enums.QuantumTimeSlot.TimeSlotType.RESERVATION and time_slot.project_id == project_id:
+                if (time_slot.slot_type == enums.QuantumTimeSlot.TimeSlotType.OPEN_SWIM) or (time_slot.slot_type == enums.QuantumTimeSlot.TimeSlotType.RESERVATION and time_slot.project_id == project_id):
                     available_processors.append(processor_name)
-                    break
     return available_processors
diff --git a/recirq/engine_utils_test.py b/recirq/engine_utils_test.py
index 6e90cd1c..774a2def 100644
--- a/recirq/engine_utils_test.py
+++ b/recirq/engine_utils_test.py
@@ -120,7 +120,7 @@ def test_sampler_by_name():
 @patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient')
 @patch('recirq.engine_utils._get_current_time')
 @patch('cirq.google.engine.EngineProcessor.get_schedule')
-def test_get_available_processors_open_swim_1(schedule_mock, time_mock, engine_mock):
+def test_get_available_processors_open_swim_in_time_window(schedule_mock, time_mock, engine_mock):
     os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project'
     schedule_mock.return_value = [
         qtypes.QuantumTimeSlot(
@@ -136,7 +136,7 @@ def test_get_available_processors_open_swim_1(schedule_mock, time_mock, engine_m
 @patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient')
 @patch('recirq.engine_utils._get_current_time')
 @patch('cirq.google.engine.EngineProcessor.get_schedule')
-def test_get_available_processors_open_swim_2(schedule_mock, time_mock, engine_mock):
+def test_get_available_processors_open_swim_outside_window(schedule_mock, time_mock, engine_mock):
     os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project'
     schedule_mock.return_value = [
         qtypes.QuantumTimeSlot(
@@ -148,3 +148,21 @@ def test_get_available_processors_open_swim_2(schedule_mock, time_mock, engine_m
     time_mock.return_value = datetime.fromtimestamp(700)
     assert recirq.get_available_processors(['Sycamore23']) == []
 
+
+@patch('cirq.google.engine.engine_client.quantum.QuantumEngineServiceClient')
+@patch('recirq.engine_utils._get_current_time')
+@patch('cirq.google.engine.EngineProcessor.get_schedule')
+def test_get_available_processors_current_project_reservation(schedule_mock, time_mock, engine_mock):
+    os.environ['GOOGLE_CLOUD_PROJECT'] = 'some_project'
+    schedule_mock.return_value = [
+        qtypes.QuantumTimeSlot(
+            processor_name='Sycamore23',
+            start_time=Timestamp(seconds=100),
+            end_time=Timestamp(seconds=500),
+            slot_type=enums.QuantumTimeSlot.TimeSlotType.RESERVATION,
+            reservation_config=qtypes.QuantumTimeSlot.ReservationConfig(project_id='some_project'),
+        )
+    ]
+    time_mock.return_value = datetime.fromtimestamp(300)
+    assert 'Sycamore23' in recirq.get_available_processors(['Sycamore23'])
+
diff --git a/recirq/quantum_chess/experiments/batch_moves.py b/recirq/quantum_chess/experiments/batch_moves.py
index 4a8435a6..ea595029 100644
--- a/recirq/quantum_chess/experiments/batch_moves.py
+++ b/recirq/quantum_chess/experiments/batch_moves.py
@@ -11,8 +11,7 @@
 by recirq.quantum_chess.move (see also interactive_board for
 examples).
 
-PROCESSOR_NAME is a processor name from engine_utils.py.
-Defaults to a 54 qubit sycamore noiseless simulator.
+PROCESSOR_NAME is a processor name from engine_utils.py. If empty, it will find an available quantum processor to run on. Default to 54 qubit sycamore noiseless simulator if none are available.
 
 FEN is a initial position in chess FEN notation. Optional.
 Default is the normal classical chess starting position.