def test_set_job_labels(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_job.return_value = qtypes.QuantumJob(
labels={
'color': 'red',
'weather': 'sun',
'run': '1'
},
label_fingerprint='hash')
result = qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
grpc_client.update_quantum_job.return_value = result
client = EngineClient()
labels = {'hello': 'world', 'color': 'blue', 'run': '1'}
assert client.set_job_labels('proj', 'prog', 'job0', labels) == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0',
labels=labels,
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
assert client.set_job_labels('proj', 'prog', 'job0', {}) == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0',
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
def test_cancel_job(client_constructor):
grpc_client = setup_mock_(client_constructor)
client = EngineClient()
assert not client.cancel_job('proj', 'prog', 'job0')
assert grpc_client.cancel_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0', )
def test_list_time_slots(client_constructor):
grpc_client = setup_mock_(client_constructor)
results = [
qtypes.QuantumTimeSlot(
processor_name='potofgold',
start_time=Timestamp(seconds=1000020000),
end_time=Timestamp(seconds=1000040000),
slot_type=qenums.QuantumTimeSlot.TimeSlotType.MAINTENANCE,
maintenance_config=qtypes.QuantumTimeSlot.MaintenanceConfig(
title='Testing',
description='Testing some new configuration.',
),
),
qtypes.QuantumTimeSlot(
processor_name='potofgold',
start_time=Timestamp(seconds=1000010000),
end_time=Timestamp(seconds=1000020000),
slot_type=qenums.QuantumTimeSlot.TimeSlotType.RESERVATION,
reservation_config=qtypes.QuantumTimeSlot.ReservationConfig(
project_id='super_secret_quantum'),
)
]
grpc_client.list_quantum_time_slots.return_value = results
client = EngineClient()
assert (client.list_time_slots('proj', 'processor0') == results)
def test_update_reservation(client_constructor):
grpc_client = setup_mock_(client_constructor)
name = 'projects/proj/processors/processor0/reservations/papar-party-44'
result = qtypes.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)
kwargs = grpc_client.update_quantum_reservation.call_args[1]
assert kwargs == {
'name':
name,
'quantum_reservation':
result,
'update_mask':
FieldMask(paths=['start_time', 'end_time', 'whitelisted_users'])
}
def test_get_reservation_exception(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_reservation.side_effect = exceptions.BadRequest(
'boom')
client = EngineClient()
with pytest.raises(EngineException, match='boom'):
client.get_reservation('proj', 'processor0', 'goog')
def test_api_doesnt_retry_not_found_errors(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_program.side_effect = exceptions.NotFound(
'not found')
client = EngineClient()
with pytest.raises(EngineException, match='not found'):
client.get_program('proj', 'prog', False)
assert grpc_client.get_quantum_program.call_count == 1
def test_api_retry_5xx_errors(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_program.side_effect = exceptions.ServiceUnavailable(
'internal error')
client = EngineClient(max_retry_delay_seconds=0.3)
with pytest.raises(TimeoutError,
match='Reached max retry attempts.*internal error'):
client.get_program('proj', 'prog', False)
assert grpc_client.get_quantum_program.call_count == 3
def test_get_current_calibration_does_not_exist(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_calibration.side_effect = exceptions.NotFound(
'not found')
client = EngineClient()
assert client.get_current_calibration('proj', 'processor0') is None
assert grpc_client.get_quantum_calibration.call_args[0] == (
'projects/proj/processors/processor0/calibrations/current', )
def test_get_processor(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumProcessor(name='projects/proj/processors/processor0')
grpc_client.get_quantum_processor.return_value = result
client = EngineClient()
assert client.get_processor('proj', 'processor0') == result
assert grpc_client.get_quantum_processor.call_args[0] == (
'projects/proj/processors/processor0',)
def test_delete_program(client_constructor):
grpc_client = setup_mock_(client_constructor)
client = EngineClient()
assert not client.delete_program('proj', 'prog')
assert grpc_client.delete_quantum_program.call_args[0] == (
'projects/proj/programs/prog', False)
assert not client.delete_program('proj', 'prog', delete_jobs=True)
assert grpc_client.delete_quantum_program.call_args[0] == (
'projects/proj/programs/prog', True)
def test_list_program_filters(client_constructor, expected_filter,
created_before, created_after, labels):
grpc_client = setup_mock_(client_constructor)
client = EngineClient()
client.list_programs(project_id='proj',
created_before=created_before,
created_after=created_after,
has_labels=labels)
assert grpc_client.list_quantum_programs.call_args[1] == {
'filter_': expected_filter,
}
def test_job_results(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumResult(
parent='projects/proj/programs/prog/jobs/job0')
grpc_client.get_quantum_result.return_value = result
client = EngineClient()
assert client.get_job_results('proj', 'prog', 'job0') == result
assert grpc_client.get_quantum_result.call_args[0] == (
'projects/proj/programs/prog/jobs/job0', )
def test_get_current_calibration(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumCalibration(
name='projects/proj/processors/processor0/calibrations/123456')
grpc_client.get_quantum_calibration.return_value = result
client = EngineClient()
assert client.get_current_calibration('proj', 'processor0') == result
assert grpc_client.get_quantum_calibration.call_args[0] == (
'projects/proj/processors/processor0/calibrations/current', )
def test_list_processors(client_constructor):
grpc_client = setup_mock_(client_constructor)
results = [
qtypes.QuantumProcessor(name='projects/proj/processor/processor0'),
qtypes.QuantumProcessor(name='projects/proj/processor/processor1')
]
grpc_client.list_quantum_processors.return_value = results
client = EngineClient()
assert client.list_processors('proj') == results
assert grpc_client.list_quantum_processors.call_args[0] == (
'projects/proj', '')
def test_get_reservation_not_found(client_constructor):
grpc_client = setup_mock_(client_constructor)
name = 'projects/proj/processors/processor0/reservations/papar-party-44'
grpc_client.get_quantum_reservation.side_effect = exceptions.NotFound(
'not found')
client = EngineClient()
assert (client.get_reservation('proj', 'processor0',
'papar-party-44') == None)
kwargs = grpc_client.get_quantum_reservation.call_args[1]
assert kwargs == {
'name': name,
}
def test_get_program(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumProgram(name='projects/proj/programs/prog')
grpc_client.get_quantum_program.return_value = result
client = EngineClient()
assert client.get_program('proj', 'prog', False) == result
assert grpc_client.get_quantum_program.call_args[0] == (
'projects/proj/programs/prog', False)
assert client.get_program('proj', 'prog', True) == result
assert grpc_client.get_quantum_program.call_args[0] == (
'projects/proj/programs/prog', True)
def test_api_retry_times(client_constructor, mock_time):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_program.side_effect = exceptions.ServiceUnavailable(
'internal error')
client = EngineClient(max_retry_delay_seconds=0.3)
with pytest.raises(TimeoutError,
match='Reached max retry attempts.*internal error'):
client.get_program('proj', 'prog', False)
assert grpc_client.get_quantum_program.call_count == 3
assert len(mock_time.call_args_list) == 2
assert all(x == y
for (x,
_), y in zip(mock_time.call_args_list, [(0.1, ), (0.2, )]))
def test_list_calibrations(client_constructor):
grpc_client = setup_mock_(client_constructor)
results = [
qtypes.QuantumCalibration(
name='projects/proj/processor/processor0/calibrations/123456'),
qtypes.QuantumCalibration(
name='projects/proj/processor/processor1/calibrations/224466')
]
grpc_client.list_quantum_calibrations.return_value = results
client = EngineClient()
assert client.list_calibrations('proj', 'processor0') == results
assert grpc_client.list_quantum_calibrations.call_args[0] == (
'projects/proj/processors/processor0', )
def test_list_program(client_constructor):
grpc_client = setup_mock_(client_constructor)
results = [
qtypes.QuantumProgram(name='projects/proj/programs/prog1'),
qtypes.QuantumProgram(name='projects/proj/programs/prog2')
]
grpc_client.list_quantum_programs.return_value = results
client = EngineClient()
assert client.list_programs(project_id='proj') == results
assert grpc_client.list_quantum_programs.call_args[0] == ('projects/proj',)
assert grpc_client.list_quantum_programs.call_args[1] == {
'filter_': '',
}
def test_add_program_labels(client_constructor):
grpc_client = setup_mock_(client_constructor)
existing = qtypes.QuantumProgram(labels={
'color': 'red',
'weather': 'sun',
'run': '1'
},
label_fingerprint='hash')
grpc_client.get_quantum_program.return_value = existing
result = qtypes.QuantumProgram(name='projects/proj/programs/prog')
grpc_client.update_quantum_program.return_value = result
client = EngineClient()
assert client.add_program_labels('proj', 'prog',
{'color': 'red'}) == existing
assert grpc_client.update_quantum_program.call_count == 0
assert client.add_program_labels('proj', 'prog',
{'hello': 'world'}) == result
assert grpc_client.update_quantum_program.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumProgram(name='projects/proj/programs/prog',
labels={
'color': 'red',
'weather': 'sun',
'run': '1',
'hello': 'world'
},
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
assert client.add_program_labels('proj', 'prog', {
'hello': 'world',
'color': 'blue'
}) == result
assert grpc_client.update_quantum_program.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumProgram(name='projects/proj/programs/prog',
labels={
'color': 'blue',
'weather': 'sun',
'run': '1',
'hello': 'world'
},
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
def test_get_reservation(client_constructor):
grpc_client = setup_mock_(client_constructor)
name = 'projects/proj/processors/processor0/reservations/papar-party-44'
result = qtypes.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)
kwargs = grpc_client.get_quantum_reservation.call_args[1]
assert kwargs == {
'name': name,
}
def test_set_job_description(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
grpc_client.update_quantum_job.return_value = result
client = EngineClient()
assert client.set_job_description('proj', 'prog', 'job0', 'A job') == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0',
description='A job'),
qtypes.field_mask_pb2.FieldMask(paths=['description']))
assert client.set_job_description('proj', 'prog', 'job0', '') == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0'),
qtypes.field_mask_pb2.FieldMask(paths=['description']))
def test_list_reservation(client_constructor):
grpc_client = setup_mock_(client_constructor)
name = 'projects/proj/processors/processor0/reservations/papar-party-44'
results = [
qtypes.QuantumReservation(
name=name,
start_time=Timestamp(seconds=1000000000),
end_time=Timestamp(seconds=1000002000),
whitelisted_users=['*****@*****.**'],
),
qtypes.QuantumReservation(
name=name,
start_time=Timestamp(seconds=1200000000),
end_time=Timestamp(seconds=1200002000),
whitelisted_users=['*****@*****.**'],
),
]
grpc_client.list_quantum_reservations.return_value = results
client = EngineClient()
assert (client.list_reservations('proj', 'processor0') == results)
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 = qtypes.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)
kwargs = grpc_client.update_quantum_reservation.call_args[1]
assert kwargs == {
'name': name,
'quantum_reservation': result,
'update_mask': FieldMask(paths=['whitelisted_users'])
}
def test_create_program(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumProgram(name='projects/proj/programs/prog')
grpc_client.create_quantum_program.return_value = result
code = qtypes.any_pb2.Any()
labels = {'hello': 'world'}
client = EngineClient()
assert client.create_program('proj', 'prog', code, 'A program',
labels) == ('prog', result)
assert grpc_client.create_quantum_program.call_args[0] == (
'projects/proj',
qtypes.QuantumProgram(name='projects/proj/programs/prog',
code=code,
description='A program',
labels=labels),
False,
)
assert client.create_program('proj', 'prog', code,
'A program') == ('prog', result)
assert grpc_client.create_quantum_program.call_args[0] == (
'projects/proj',
qtypes.QuantumProgram(name='projects/proj/programs/prog',
code=code,
description='A program'),
False,
)
assert client.create_program('proj', 'prog', code,
labels=labels) == ('prog', result)
assert grpc_client.create_quantum_program.call_args[0] == (
'projects/proj',
qtypes.QuantumProgram(name='projects/proj/programs/prog',
code=code,
labels=labels),
False,
)
assert client.create_program('proj', 'prog', code) == ('prog', result)
assert grpc_client.create_quantum_program.call_args[0] == (
'projects/proj',
qtypes.QuantumProgram(name='projects/proj/programs/prog', code=code),
False,
)
assert client.create_program('proj', program_id=None,
code=code) == ('prog', result)
assert grpc_client.create_quantum_program.call_args[0] == (
'projects/proj',
qtypes.QuantumProgram(code=code),
False,
)
def test_remove_job_labels(client_constructor):
grpc_client = setup_mock_(client_constructor)
existing = qtypes.QuantumJob(labels={
'color': 'red',
'weather': 'sun',
'run': '1'
},
label_fingerprint='hash')
grpc_client.get_quantum_job.return_value = existing
result = qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
grpc_client.update_quantum_job.return_value = result
client = EngineClient()
assert client.remove_job_labels('proj', 'prog', 'job0',
['other']) == existing
assert grpc_client.update_quantum_program.call_count == 0
assert client.remove_job_labels('proj', 'prog', 'job0',
['hello', 'weather']) == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0',
labels={
'color': 'red',
'run': '1',
},
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
assert client.remove_job_labels('proj', 'prog', 'job0',
['color', 'weather', 'run']) == result
assert grpc_client.update_quantum_job.call_args[0] == (
'projects/proj/programs/prog/jobs/job0',
qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0',
label_fingerprint='hash'),
qtypes.field_mask_pb2.FieldMask(paths=['labels']))
def test_create_reservation(client_constructor):
grpc_client = setup_mock_(client_constructor)
start = datetime.datetime.fromtimestamp(1000000000)
end = datetime.datetime.fromtimestamp(1000003600)
users = ['*****@*****.**']
result = qtypes.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
kwargs = grpc_client.create_quantum_reservation.call_args[1]
# The outgoing argument will not have the resource name
result.name = ''
assert kwargs == {
'parent': 'projects/proj/processors/processor0',
'quantum_reservation': result
}
def test_create_job(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = qtypes.QuantumJob(name='projects/proj/programs/prog/jobs/job0')
grpc_client.create_quantum_job.return_value = result
run_context = qtypes.any_pb2.Any()
labels = {'hello': 'world'}
client = EngineClient()
assert client.create_job('proj', 'prog', 'job0', ['processor0'],
run_context, 10, 'A job',
labels) == ('job0', result)
assert grpc_client.create_quantum_job.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
run_context=run_context,
scheduling_config=qtypes.SchedulingConfig(
priority=10,
processor_selector=qtypes.SchedulingConfig.ProcessorSelector(
processor_names=['projects/proj/processors/processor0'])),
description='A job',
labels=labels), False)
assert client.create_job(
'proj',
'prog',
'job0',
['processor0'],
run_context,
10,
'A job',
) == ('job0', result)
assert grpc_client.create_quantum_job.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
run_context=run_context,
scheduling_config=qtypes.SchedulingConfig(
priority=10,
processor_selector=qtypes.SchedulingConfig.ProcessorSelector(
processor_names=['projects/proj/processors/processor0'])),
description='A job'), False)
assert client.create_job('proj',
'prog',
'job0', ['processor0'],
run_context,
10,
labels=labels) == ('job0', result)
assert grpc_client.create_quantum_job.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
run_context=run_context,
scheduling_config=qtypes.SchedulingConfig(
priority=10,
processor_selector=qtypes.SchedulingConfig.ProcessorSelector(
processor_names=['projects/proj/processors/processor0'])),
labels=labels), False)
assert client.create_job('proj', 'prog', 'job0', ['processor0'],
run_context, 10) == ('job0', result)
assert grpc_client.create_quantum_job.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
run_context=run_context,
scheduling_config=qtypes.SchedulingConfig(
priority=10,
processor_selector=qtypes.SchedulingConfig.ProcessorSelector(
processor_names=['projects/proj/processors/processor0'])),
), False)
assert client.create_job('proj',
'prog',
job_id=None,
processor_ids=['processor0'],
run_context=run_context,
priority=10) == ('job0', result)
assert grpc_client.create_quantum_job.call_args[0] == (
'projects/proj/programs/prog',
qtypes.QuantumJob(
run_context=run_context,
scheduling_config=qtypes.SchedulingConfig(
priority=10,
processor_selector=qtypes.SchedulingConfig.ProcessorSelector(
processor_names=['projects/proj/processors/processor0'])),
), False)
with pytest.raises(ValueError,
match='priority must be between 0 and 1000'):
client.create_job('proj',
'prog',
job_id=None,
processor_ids=['processor0'],
run_context=run_context,
priority=5000)
def test_list_program_filters_invalid_type(client_constructor):
with pytest.raises(ValueError, match=""):
EngineClient().list_programs(project_id='proj',
created_before="Unsupported date/time")