def test_list_reservation(list_reservations): name = 'projects/proj/processors/p0/reservations/rid' results = [ quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000003600), whitelisted_users=['*****@*****.**'], ), quantum.QuantumReservation( name=name + '2', start_time=Timestamp(seconds=1000003600), end_time=Timestamp(seconds=1000007200), whitelisted_users=['*****@*****.**'], ), ] list_reservations.return_value = results processor = cg.EngineProcessor('proj', 'p0', EngineContext()) assert ( processor.list_reservations( datetime.datetime.fromtimestamp(1000000000), datetime.datetime.fromtimestamp(1000010000) ) == results ) list_reservations.assert_called_once_with( 'proj', 'p0', 'start_time < 1000010000 AND end_time > 1000000000' )
def test_update_reservation(client_constructor): grpc_client = setup_mock_(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000001000), end_time=Timestamp(seconds=1000002000), whitelisted_users=['*****@*****.**'], ) grpc_client.update_quantum_reservation.return_value = result client = EngineClient() assert ( client.update_reservation( 'proj', 'processor0', 'papar-party-44', start=datetime.datetime.fromtimestamp(1000001000), end=datetime.datetime.fromtimestamp(1000002000), whitelisted_users=['*****@*****.**'], ) == result ) grpc_client.update_quantum_reservation.assert_called_with( quantum.UpdateQuantumReservationRequest( name=name, quantum_reservation=result, update_mask=FieldMask(paths=['start_time', 'end_time', 'whitelisted_users']), ) )
def create_reservation( self, project_id: str, processor_id: str, start: datetime.datetime, end: datetime.datetime, whitelisted_users: Optional[List[str]] = None, ): """Creates a quantum reservation and returns the created object. Args: project_id: A project_id of the parent Google Cloud Project. processor_id: The processor unique identifier. reservation_id: Unique ID of the reservation in the parent project, or None if the engine should generate an id start: the starting time of the reservation as a datetime object end: the ending time of the reservation as a datetime object whitelisted_users: a list of emails that can use the reservation. """ parent = _processor_name_from_ids(project_id, processor_id) reservation = quantum.QuantumReservation( name='', start_time=Timestamp(seconds=int(start.timestamp())), end_time=Timestamp(seconds=int(end.timestamp())), ) if whitelisted_users: reservation.whitelisted_users.extend(whitelisted_users) request = quantum.CreateQuantumReservationRequest( parent=parent, quantum_reservation=reservation) return self._make_request(self.grpc_client.create_quantum_reservation, request)
def create_reservation( self, start_time: datetime.datetime, end_time: datetime.datetime, whitelisted_users: Optional[List[str]] = None, ) -> quantum.QuantumReservation: """Creates a reservation on this processor. Args: start_time: the starting date/time of the reservation. end_time: the ending date/time of the reservation. whitelisted_users: a list of emails that are allowed to send programs during this reservation (in addition to users with permission "quantum.reservations.use" on the project). Raises: ValueError: if start_time is after end_time. """ if end_time < start_time: raise ValueError( 'End time of reservation must be after the start time') reservation_id = self._create_id() new_reservation = quantum.QuantumReservation( name=reservation_id, start_time=Timestamp(seconds=int(start_time.timestamp())), end_time=Timestamp(seconds=int(end_time.timestamp())), whitelisted_users=whitelisted_users, ) time_slot = self._reservation_to_time_slot(new_reservation) if not self._is_available(time_slot): raise ValueError('Time slot is not available for reservations') self._reservations[reservation_id] = new_reservation self._insert_reservation_into(time_slot) return new_reservation
def test_get_reservation(get_reservation): name = 'projects/proj/processors/p0/reservations/rid' result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000003600), whitelisted_users=['*****@*****.**'], ) get_reservation.return_value = result processor = cg.EngineProcessor('proj', 'p0', EngineContext()) assert processor.get_reservation('rid') == result get_reservation.assert_called_once_with('proj', 'p0', 'rid')
def test_list_reservation(client_constructor): grpc_client = setup_mock_(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' results = [ quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000002000), whitelisted_users=['*****@*****.**'], ), quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1200000000), end_time=Timestamp(seconds=1200002000), whitelisted_users=['*****@*****.**'], ), ] grpc_client.list_quantum_reservations.return_value = Pager(results) client = EngineClient() assert client.list_reservations('proj', 'processor0') == results
async def update_reservation_async( self, project_id: str, processor_id: str, reservation_id: str, start: Optional[datetime.datetime] = None, end: Optional[datetime.datetime] = None, whitelisted_users: Optional[List[str]] = None, ): """Updates a quantum reservation. This will update a quantum reservation's starting time, ending time, and list of whitelisted users. If any field is not filled, it will not be updated. Args: project_id: A project_id of the parent Google Cloud Project. processor_id: The processor unique identifier. reservation_id: Unique ID of the reservation in the parent project, start: the new starting time of the reservation as a datetime object end: the new ending time of the reservation as a datetime object whitelisted_users: a list of emails that can use the reservation. The empty list, [], will clear the whitelisted_users while None will leave the value unchanged. """ name = ( _reservation_name_from_ids(project_id, processor_id, reservation_id) if reservation_id else '' ) reservation = quantum.QuantumReservation(name=name) paths = [] if start: reservation.start_time = start paths.append('start_time') if end: reservation.end_time = end paths.append('end_time') if whitelisted_users is not None: reservation.whitelisted_users.extend(whitelisted_users) paths.append('whitelisted_users') request = quantum.UpdateQuantumReservationRequest( name=name, quantum_reservation=reservation, update_mask=field_mask_pb2.FieldMask(paths=paths), ) return await self._send_request_async(self.grpc_client.update_quantum_reservation, request)
def test_get_reservation(client_constructor): grpc_client = setup_mock_(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000002000), whitelisted_users=['*****@*****.**'], ) grpc_client.get_quantum_reservation.return_value = result client = EngineClient() assert client.get_reservation('proj', 'processor0', 'papar-party-44') == result grpc_client.get_quantum_reservation.assert_called_with( quantum.GetQuantumReservationRequest(name=name) )
def test_update_reservation(update_reservation): name = 'projects/proj/processors/p0/reservations/rid' result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000003600), whitelisted_users=['*****@*****.**'], ) start = datetime.datetime.fromtimestamp(1000000000) end = datetime.datetime.fromtimestamp(1000003600) update_reservation.return_value = result processor = cg.EngineProcessor('proj', 'p0', EngineContext()) assert processor.update_reservation('rid', start, end, ['*****@*****.**']) == result update_reservation.assert_called_once_with( 'proj', 'p0', 'rid', start=start, end=end, whitelisted_users=['*****@*****.**'] )
def test_update_reservation_remove_all_users(client_constructor): grpc_client = setup_mock_(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation(name=name, whitelisted_users=[]) grpc_client.update_quantum_reservation.return_value = result client = EngineClient() assert ( client.update_reservation('proj', 'processor0', 'papar-party-44', whitelisted_users=[]) == result ) grpc_client.update_quantum_reservation.assert_called_with( quantum.UpdateQuantumReservationRequest( name=name, quantum_reservation=result, update_mask=FieldMask(paths=['whitelisted_users']), ) )
def test_remove_reservation_cancel(cancel_reservation, get_reservation): name = 'projects/proj/processors/p0/reservations/rid' now = int(datetime.datetime.now().timestamp()) result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=now + 10), end_time=Timestamp(seconds=now + 3610), whitelisted_users=['*****@*****.**'], ) get_reservation.return_value = result cancel_reservation.return_value = result processor = cg.EngineProcessor( 'proj', 'p0', EngineContext(), quantum.QuantumProcessor(schedule_frozen_period=Duration(seconds=10000)), ) assert processor.remove_reservation('rid') == result cancel_reservation.assert_called_once_with('proj', 'p0', 'rid')
def test_remove_reservation_failures(get_reservation, get_processor): name = 'projects/proj/processors/p0/reservations/rid' now = int(datetime.datetime.now().timestamp()) result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=now + 10), end_time=Timestamp(seconds=now + 3610), whitelisted_users=['*****@*****.**'], ) get_reservation.return_value = result get_processor.return_value = None # no processor processor = cg.EngineProcessor('proj', 'p0', EngineContext()) with pytest.raises(ValueError): processor.remove_reservation('rid') # No freeze period defined processor = cg.EngineProcessor('proj', 'p0', EngineContext(), quantum.QuantumProcessor()) with pytest.raises(ValueError): processor.remove_reservation('rid')
def test_create_reservation(create_reservation): name = 'projects/proj/processors/p0/reservations/psherman-wallaby-way' result = quantum.QuantumReservation( name=name, start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000003600), whitelisted_users=['*****@*****.**'], ) create_reservation.return_value = result processor = cg.EngineProcessor('proj', 'p0', EngineContext()) assert processor.create_reservation( datetime.datetime.fromtimestamp(1000000000), datetime.datetime.fromtimestamp(1000003600), ['*****@*****.**'], ) create_reservation.assert_called_once_with( 'proj', 'p0', datetime.datetime.fromtimestamp(1000000000), datetime.datetime.fromtimestamp(1000003600), ['*****@*****.**'], )
def test_create_reservation(client_constructor): grpc_client = setup_mock_(client_constructor) start = datetime.datetime.fromtimestamp(1000000000) end = datetime.datetime.fromtimestamp(1000003600) users = ['*****@*****.**'] result = quantum.QuantumReservation( name='projects/proj/processors/processor0/reservations/papar-party-44', start_time=Timestamp(seconds=1000000000), end_time=Timestamp(seconds=1000003600), whitelisted_users=users, ) grpc_client.create_quantum_reservation.return_value = result client = EngineClient() assert client.create_reservation('proj', 'processor0', start, end, users) == result assert grpc_client.create_quantum_reservation.call_count == 1 # The outgoing argument will not have the resource name result.name = '' grpc_client.create_quantum_reservation.assert_called_with( quantum.CreateQuantumReservationRequest( parent='projects/proj/processors/processor0', quantum_reservation=result ) )