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
)
)
