def test_set_job_labels(client_constructor):
grpc_client = setup_mock_(client_constructor)
grpc_client.get_quantum_job.return_value = quantum.QuantumJob(
labels={'color': 'red', 'weather': 'sun', 'run': '1'}, label_fingerprint='hash'
)
result = quantum.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
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
labels=labels,
label_fingerprint='hash',
),
update_mask=FieldMask(paths=['labels']),
)
)
assert client.set_job_labels('proj', 'prog', 'job0', {}) == result
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(
name='projects/proj/programs/prog/jobs/job0', label_fingerprint='hash'
),
update_mask=FieldMask(paths=['labels']),
)
)
def test_set_job_description(client_constructor):
grpc_client = setup_mock_(client_constructor)
result = quantum.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
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(
name='projects/proj/programs/prog/jobs/job0', description='A job'
),
update_mask=FieldMask(paths=['description']),
)
)
assert client.set_job_description('proj', 'prog', 'job0', '') == result
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(name='projects/proj/programs/prog/jobs/job0'),
update_mask=FieldMask(paths=['description']),
)
)
def test_remove_job_labels(client_constructor):
grpc_client = setup_mock_(client_constructor)
existing = quantum.QuantumJob(
labels={'color': 'red', 'weather': 'sun', 'run': '1'}, label_fingerprint='hash'
)
grpc_client.get_quantum_job.return_value = existing
result = quantum.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
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(
name='projects/proj/programs/prog/jobs/job0',
labels={'color': 'red', 'run': '1'},
label_fingerprint='hash',
),
update_mask=FieldMask(paths=['labels']),
)
)
assert client.remove_job_labels('proj', 'prog', 'job0', ['color', 'weather', 'run']) == result
grpc_client.update_quantum_job.assert_called_with(
quantum.UpdateQuantumJobRequest(
name='projects/proj/programs/prog/jobs/job0',
quantum_job=quantum.QuantumJob(
name='projects/proj/programs/prog/jobs/job0', label_fingerprint='hash'
),
update_mask=FieldMask(paths=['labels']),
)
)
def _set_job_labels(
self,
project_id: str,
program_id: str,
job_id: str,
labels: Dict[str, str],
fingerprint: str,
) -> quantum.QuantumJob:
job_resource_name = _job_name_from_ids(project_id, program_id, job_id)
request = quantum.UpdateQuantumJobRequest(
name=job_resource_name,
quantum_job=quantum.QuantumJob(name=job_resource_name,
labels=labels,
label_fingerprint=fingerprint),
update_mask=field_mask_pb2.FieldMask(paths=['labels']),
)
return self._make_request(self.grpc_client.update_quantum_job, request)
async def set_job_description_async(
self, project_id: str, program_id: str, job_id: str, description: str
) -> quantum.QuantumJob:
"""Sets the description for a previously created quantum job.
Args:
project_id: A project_id of the parent Google Cloud Project.
program_id: Unique ID of the program within the parent project.
job_id: Unique ID of the job within the parent program.
description: The new job description.
Returns:
The updated quantum job.
"""
job_resource_name = _job_name_from_ids(project_id, program_id, job_id)
request = quantum.UpdateQuantumJobRequest(
name=job_resource_name,
quantum_job=quantum.QuantumJob(name=job_resource_name, description=description),
update_mask=field_mask_pb2.FieldMask(paths=['description']),
)
return await self._send_request_async(self.grpc_client.update_quantum_job, request)