From 778e88626de461bf4bb410a54b8de9191f1619d2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Doreau?= Date: Wed, 8 Jan 2025 17:54:31 +0200 Subject: [PATCH] Refactor tests to use pytest fixtures for `default_run_context` and `default_engine_client` - Updated all test cases to use these fixtures directly for consistency and clarity. --- .../cirq_google/engine/engine_client_test.py | 428 +++++++++--------- 1 file changed, 211 insertions(+), 217 deletions(-) diff --git a/cirq-google/cirq_google/engine/engine_client_test.py b/cirq-google/cirq_google/engine/engine_client_test.py index b741184eaa3..0bd4c0f8c9c 100644 --- a/cirq-google/cirq_google/engine/engine_client_test.py +++ b/cirq-google/cirq_google/engine/engine_client_test.py @@ -31,14 +31,17 @@ # JOB_PATH represents the path to a specific job. JOB_PATH = 'projects/proj/programs/prog/jobs/job0' + @pytest.fixture def default_run_context(): return any_pb2.Any() + @pytest.fixture def default_engine_client(): return EngineClient() + def _setup_client_mock(client_constructor): grpc_client = mock.AsyncMock() client_constructor.return_value = grpc_client @@ -52,7 +55,7 @@ def _setup_stream_manager_mock(manager_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_create_program(client_constructor): +def test_create_program(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumProgram(name='projects/proj/programs/prog') @@ -60,9 +63,11 @@ def test_create_program(client_constructor): code = any_pb2.Any() labels = {'hello': 'world'} - client = EngineClient() - assert client.create_program('proj', 'prog', code, 'A program', labels) == ('prog', result) + assert default_engine_client.create_program('proj', 'prog', code, 'A program', labels) == ( + 'prog', + result, + ) grpc_client.create_quantum_program.assert_called_with( quantum.CreateQuantumProgramRequest( parent='projects/proj', @@ -75,7 +80,10 @@ def test_create_program(client_constructor): ) ) - assert client.create_program('proj', 'prog', code, 'A program') == ('prog', result) + assert default_engine_client.create_program('proj', 'prog', code, 'A program') == ( + 'prog', + result, + ) grpc_client.create_quantum_program.assert_called_with( quantum.CreateQuantumProgramRequest( parent='projects/proj', @@ -85,7 +93,10 @@ def test_create_program(client_constructor): ) ) - assert client.create_program('proj', 'prog', code, labels=labels) == ('prog', result) + assert default_engine_client.create_program('proj', 'prog', code, labels=labels) == ( + 'prog', + result, + ) grpc_client.create_quantum_program.assert_called_with( quantum.CreateQuantumProgramRequest( parent='projects/proj', @@ -95,7 +106,7 @@ def test_create_program(client_constructor): ) ) - assert client.create_program('proj', 'prog', code) == ('prog', result) + assert default_engine_client.create_program('proj', 'prog', code) == ('prog', result) grpc_client.create_quantum_program.assert_called_with( quantum.CreateQuantumProgramRequest( parent='projects/proj', @@ -103,7 +114,10 @@ def test_create_program(client_constructor): ) ) - assert client.create_program('proj', program_id=None, code=code) == ('prog', result) + assert default_engine_client.create_program('proj', program_id=None, code=code) == ( + 'prog', + result, + ) grpc_client.create_quantum_program.assert_called_with( quantum.CreateQuantumProgramRequest( parent='projects/proj', quantum_program=quantum.QuantumProgram(code=code) @@ -112,26 +126,25 @@ def test_create_program(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_program(client_constructor): +def test_get_program(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.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 default_engine_client.get_program('proj', 'prog', False) == result grpc_client.get_quantum_program.assert_called_with( quantum.GetQuantumProgramRequest(name='projects/proj/programs/prog', return_code=False) ) - assert client.get_program('proj', 'prog', True) == result + assert default_engine_client.get_program('proj', 'prog', True) == result grpc_client.get_quantum_program.assert_called_with( quantum.GetQuantumProgramRequest(name='projects/proj/programs/prog', return_code=True) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_program(client_constructor): +def test_list_program(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) results = [ @@ -140,8 +153,7 @@ def test_list_program(client_constructor): ] grpc_client.list_quantum_programs.return_value = results - client = EngineClient() - assert client.list_programs(project_id='proj') == results + assert default_engine_client.list_programs(project_id='proj') == results grpc_client.list_quantum_programs.assert_called_with( quantum.ListQuantumProgramsRequest(parent='projects/proj', filter='') ) @@ -178,11 +190,15 @@ def test_list_program(client_constructor): ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) def test_list_program_filters( - client_constructor, expected_filter, created_before, created_after, labels + client_constructor, + expected_filter, + created_before, + created_after, + labels, + default_engine_client, ): grpc_client = _setup_client_mock(client_constructor) - client = EngineClient() - client.list_programs( + default_engine_client.list_programs( project_id='proj', created_before=created_before, created_after=created_after, @@ -192,20 +208,21 @@ def test_list_program_filters( @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_program_filters_invalid_type(client_constructor): +def test_list_program_filters_invalid_type(client_constructor, default_engine_client): with pytest.raises(ValueError, match=""): - EngineClient().list_programs(project_id='proj', created_before="Unsupported date/time") + default_engine_client.list_programs( + project_id='proj', created_before="Unsupported date/time" + ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_set_program_description(client_constructor): +def test_set_program_description(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumProgram(name='projects/proj/programs/prog') grpc_client.update_quantum_program.return_value = result - client = EngineClient() - assert client.set_program_description('proj', 'prog', 'A program') == result + assert default_engine_client.set_program_description('proj', 'prog', 'A program') == result grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -216,7 +233,7 @@ def test_set_program_description(client_constructor): ) ) - assert client.set_program_description('proj', 'prog', '') == result + assert default_engine_client.set_program_description('proj', 'prog', '') == result grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -227,7 +244,7 @@ def test_set_program_description(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_set_program_labels(client_constructor): +def test_set_program_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) grpc_client.get_quantum_program.return_value = quantum.QuantumProgram( @@ -236,9 +253,8 @@ def test_set_program_labels(client_constructor): result = quantum.QuantumProgram(name='projects/proj/programs/prog') grpc_client.update_quantum_program.return_value = result - client = EngineClient() labels = {'hello': 'world', 'color': 'blue', 'run': '1'} - assert client.set_program_labels('proj', 'prog', labels) == result + assert default_engine_client.set_program_labels('proj', 'prog', labels) == result grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -249,7 +265,7 @@ def test_set_program_labels(client_constructor): ) ) - assert client.set_program_labels('proj', 'prog', {}) == result + assert default_engine_client.set_program_labels('proj', 'prog', {}) == result grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -262,7 +278,7 @@ def test_set_program_labels(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_add_program_labels(client_constructor): +def test_add_program_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) existing = quantum.QuantumProgram( @@ -272,11 +288,10 @@ def test_add_program_labels(client_constructor): result = quantum.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 default_engine_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 default_engine_client.add_program_labels('proj', 'prog', {'hello': 'world'}) == result grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -289,7 +304,12 @@ def test_add_program_labels(client_constructor): ) ) - assert client.add_program_labels('proj', 'prog', {'hello': 'world', 'color': 'blue'}) == result + assert ( + default_engine_client.add_program_labels( + 'proj', 'prog', {'hello': 'world', 'color': 'blue'} + ) + == result + ) grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -304,7 +324,7 @@ def test_add_program_labels(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_remove_program_labels(client_constructor): +def test_remove_program_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) existing = quantum.QuantumProgram( @@ -314,11 +334,12 @@ def test_remove_program_labels(client_constructor): result = quantum.QuantumProgram(name='projects/proj/programs/prog') grpc_client.update_quantum_program.return_value = result - client = EngineClient() - assert client.remove_program_labels('proj', 'prog', ['other']) == existing + assert default_engine_client.remove_program_labels('proj', 'prog', ['other']) == existing assert grpc_client.update_quantum_program.call_count == 0 - assert client.remove_program_labels('proj', 'prog', ['hello', 'weather']) == result + assert ( + default_engine_client.remove_program_labels('proj', 'prog', ['hello', 'weather']) == result + ) grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -331,7 +352,10 @@ def test_remove_program_labels(client_constructor): ) ) - assert client.remove_program_labels('proj', 'prog', ['color', 'weather', 'run']) == result + assert ( + default_engine_client.remove_program_labels('proj', 'prog', ['color', 'weather', 'run']) + == result + ) grpc_client.update_quantum_program.assert_called_with( quantum.UpdateQuantumProgramRequest( name='projects/proj/programs/prog', @@ -347,13 +371,12 @@ def test_remove_program_labels(client_constructor): def test_delete_program(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) - client = default_engine_client - assert not client.delete_program('proj', 'prog') + assert not default_engine_client.delete_program('proj', 'prog') grpc_client.delete_quantum_program.assert_called_with( quantum.DeleteQuantumProgramRequest(name='projects/proj/programs/prog', delete_jobs=False) ) - assert not client.delete_program('proj', 'prog', delete_jobs=True) + assert not default_engine_client.delete_program('proj', 'prog', delete_jobs=True) grpc_client.delete_quantum_program.assert_called_with( quantum.DeleteQuantumProgramRequest(name='projects/proj/programs/prog', delete_jobs=True) ) @@ -362,9 +385,7 @@ def test_delete_program(client_constructor, default_engine_client): @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) def test_create_job_with_all_parameters( - client_constructor, - default_run_context, - default_engine_client + client_constructor, default_run_context, default_engine_client ): grpc_client = _setup_client_mock(client_constructor) @@ -372,8 +393,7 @@ def test_create_job_with_all_parameters( grpc_client.create_quantum_job.return_value = result labels = {'hello': 'world'} - client = default_engine_client - assert client.create_job( + assert default_engine_client.create_job( project_id='proj', program_id='prog', job_id='job0', @@ -381,7 +401,7 @@ def test_create_job_with_all_parameters( run_context=default_run_context, priority=10, description='A job', - labels=labels + labels=labels, ) == ('job0', result) grpc_client.create_quantum_job.assert_called_with( quantum.CreateQuantumJobRequest( @@ -405,26 +425,22 @@ def test_create_job_with_all_parameters( @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_create_job_without_labels(client_constructor): +def test_create_job_without_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result run_context = any_pb2.Any() - client = EngineClient() - assert client.create_job( + assert default_engine_client.create_job( project_id='proj', program_id='prog', job_id='job0', processor_id='processor0', run_context=run_context, priority=10, - description='A job' - ) == ( - 'job0', - result, - ) + description='A job', + ) == ('job0', result) grpc_client.create_quantum_job.assert_called_with( quantum.CreateQuantumJobRequest( parent='projects/proj/programs/prog', @@ -446,7 +462,7 @@ def test_create_job_without_labels(client_constructor): @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_create_job_without_description(client_constructor): +def test_create_job_without_description(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) @@ -454,15 +470,14 @@ def test_create_job_without_description(client_constructor): run_context = any_pb2.Any() labels = {'hello': 'world'} - client = EngineClient() - assert client.create_job( + assert default_engine_client.create_job( project_id='proj', program_id='prog', job_id='job0', processor_id='processor0', run_context=run_context, priority=10, - labels=labels + labels=labels, ) == ('job0', result) grpc_client.create_quantum_job.assert_called_with( quantum.CreateQuantumJobRequest( @@ -485,15 +500,14 @@ def test_create_job_without_description(client_constructor): @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_create_job_without_job_id(client_constructor): +def test_create_job_without_job_id(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result run_context = any_pb2.Any() - client = EngineClient() - assert client.create_job( + assert default_engine_client.create_job( 'proj', 'prog', job_id=None, processor_id='processor0', run_context=run_context, priority=10 ) == ('job0', result) grpc_client.create_quantum_job.assert_called_with( @@ -516,24 +530,21 @@ def test_create_job_without_job_id(client_constructor): @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) def test_create_job_with_invalid_priority( - client_constructor, - default_run_context, - default_engine_client + client_constructor, default_run_context, default_engine_client ): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result - client = default_engine_client with pytest.raises(ValueError, match='priority must be between 0 and 1000'): - client.create_job( + default_engine_client.create_job( project_id='proj', program_id='prog', job_id=None, processor_id='processor0', run_context=default_run_context, - priority=5000 + priority=5000, ) @@ -555,15 +566,20 @@ def test_create_job_with_invalid_priority( ], ) def test_create_job_with_invalid_processor_and_device_config_arguments_throws( - client_constructor, processor_id, run_name, snapshot_id, device_config_name, error_message + client_constructor, + processor_id, + run_name, + snapshot_id, + device_config_name, + error_message, + default_engine_client, ): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result - client = EngineClient() with pytest.raises(ValueError, match=error_message): - client.create_job( + default_engine_client.create_job( project_id='proj', program_id='prog', job_id=None, @@ -581,15 +597,14 @@ def test_create_job_with_invalid_processor_and_device_config_arguments_throws( [('RUN_NAME', '', 'CONFIG_NAME'), ('', '', ''), ('', '', '')], ) def test_create_job_with_run_name_and_device_config_name_succeeds( - client_constructor, run_name, snapshot_id, device_config_name + client_constructor, run_name, snapshot_id, device_config_name, default_engine_client ): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result run_context = any_pb2.Any() - client = EngineClient() - assert client.create_job( + assert default_engine_client.create_job( project_id='proj', program_id='prog', job_id='job0', @@ -623,25 +638,23 @@ def test_create_job_with_run_name_and_device_config_name_succeeds( @mock.patch.dict(os.environ, clear='CIRQ_TESTING') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) def test_create_job_with_snapshot_id_and_config_successfully_passes_device_config_selector( - client_constructor, + client_constructor, default_engine_client, default_run_context ): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.create_quantum_job.return_value = result - run_context = any_pb2.Any() processor_id = "processor0" snapshot_id = "SNAPSHOT_ID" device_config_name = "DEVICE_CONFIG_NAME" - client = EngineClient() - client.create_job( + default_engine_client.create_job( project_id='proj', program_id='prog', job_id='job0', processor_id=processor_id, snapshot_id=snapshot_id, device_config_name=device_config_name, - run_context=run_context, + run_context=default_run_context, priority=10, ) @@ -865,7 +878,11 @@ def test_create_job_with_snapshot_id_and_config_successfully_passes_device_confi @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) @mock.patch.object(engine_stream_manager, 'StreamManager', autospec=True) def test_run_job_over_stream( - manager_constructor, client_constructor, run_job_kwargs, expected_submit_args + manager_constructor, + client_constructor, + run_job_kwargs, + expected_submit_args, + default_engine_client, ): _setup_client_mock(client_constructor) stream_manager = _setup_stream_manager_mock(manager_constructor) @@ -873,9 +890,8 @@ def test_run_job_over_stream( parent = expected_submit_args[2].name expected_future = duet.futuretools.completed_future(quantum.QuantumResult(parent=parent)) stream_manager.submit.return_value = expected_future - client = EngineClient() - actual_future = client.run_job_over_stream(**run_job_kwargs) + actual_future = default_engine_client.run_job_over_stream(**run_job_kwargs) assert actual_future == expected_future stream_manager.submit.assert_called_with(*expected_submit_args) @@ -884,11 +900,10 @@ def test_run_job_over_stream( @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) @mock.patch.object(engine_stream_manager, 'StreamManager', autospec=True) def test_run_job_over_stream_with_snapshot_id_returns_correct_future( - manager_constructor, client_constructor + manager_constructor, client_constructor, default_engine_client ): _setup_client_mock(client_constructor) stream_manager = _setup_stream_manager_mock(manager_constructor) - client = EngineClient() run_job_kwargs = ( { 'project_id': 'proj', @@ -902,13 +917,11 @@ def test_run_job_over_stream_with_snapshot_id_returns_correct_future( }, ) - expected_future = duet.futuretools.completed_future( - quantum.QuantumResult(parent=JOB_PATH) - ) + expected_future = duet.futuretools.completed_future(quantum.QuantumResult(parent=JOB_PATH)) stream_manager.submit.return_value = expected_future stream_manager.submit.return_value = expected_future - actual_future = client.run_job_over_stream(**run_job_kwargs[0]) + actual_future = default_engine_client.run_job_over_stream(**run_job_kwargs[0]) assert actual_future == expected_future @@ -916,11 +929,10 @@ def test_run_job_over_stream_with_snapshot_id_returns_correct_future( @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) @mock.patch.object(engine_stream_manager, 'StreamManager', autospec=True) def test_run_job_over_stream_with_snapshot_id_propogates_snapshot_id( - manager_constructor, client_constructor + manager_constructor, client_constructor, default_engine_client, default_run_context ): _setup_client_mock(client_constructor) stream_manager = _setup_stream_manager_mock(manager_constructor) - client = EngineClient() run_job_kwargs = ( { 'project_id': 'proj', @@ -939,7 +951,7 @@ def test_run_job_over_stream_with_snapshot_id_propogates_snapshot_id( quantum.QuantumProgram(name='projects/proj/programs/prog', code=any_pb2.Any()), quantum.QuantumJob( name=JOB_PATH, - run_context=any_pb2.Any(), + run_context=default_run_context, scheduling_config=quantum.SchedulingConfig( processor_selector=quantum.SchedulingConfig.ProcessorSelector( processor='projects/proj/processors/processor0', @@ -956,37 +968,37 @@ def test_run_job_over_stream_with_snapshot_id_propogates_snapshot_id( stream_manager.submit.return_value = expected_future stream_manager.submit.return_value = expected_future - _ = client.run_job_over_stream(**run_job_kwargs[0]) + _ = default_engine_client.run_job_over_stream(**run_job_kwargs[0]) stream_manager.submit.assert_called_with(*expected_submit_args[0]) -def test_run_job_over_stream_with_priority_out_of_bound_raises(): - client = EngineClient() +def test_run_job_over_stream_with_priority_out_of_bound_raises( + default_engine_client, default_run_context +): with pytest.raises(ValueError): - client.run_job_over_stream( + default_engine_client.run_job_over_stream( project_id='proj', program_id='prog', code=any_pb2.Any(), job_id='job0', processor_id='processor0', - run_context=any_pb2.Any(), + run_context=default_run_context, priority=9001, ) -def test_run_job_over_stream_processor_unset_raises(): - client = EngineClient() +def test_run_job_over_stream_processor_unset_raises(default_engine_client, default_run_context): with pytest.raises(ValueError, match='Must specify a processor id'): - client.run_job_over_stream( + default_engine_client.run_job_over_stream( project_id='proj', program_id='prog', code=any_pb2.Any(), job_id='job0', processor_id='', - run_context=any_pb2.Any(), + run_context=default_run_context, ) @@ -999,18 +1011,22 @@ def test_run_job_over_stream_processor_unset_raises(): ], ) def test_run_job_over_stream_invalid_device_config_raises( - run_name, snapshot_id, device_config_name, error_message + run_name, + snapshot_id, + device_config_name, + error_message, + default_engine_client, + default_run_context, ): - client = EngineClient() with pytest.raises(ValueError, match=error_message): - client.run_job_over_stream( + default_engine_client.run_job_over_stream( project_id='proj', program_id='prog', code=any_pb2.Any(), job_id='job0', processor_id='mysim', - run_context=any_pb2.Any(), + run_context=default_run_context, run_name=run_name, snapshot_id=snapshot_id, device_config_name=device_config_name, @@ -1018,48 +1034,40 @@ def test_run_job_over_stream_invalid_device_config_raises( @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_job(client_constructor): +def test_get_job(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.get_quantum_job.return_value = result - client = EngineClient() - assert client.get_job('proj', 'prog', 'job0', False) == result + assert default_engine_client.get_job('proj', 'prog', 'job0', False) == result grpc_client.get_quantum_job.assert_called_with( - quantum.GetQuantumJobRequest( - name=JOB_PATH, return_run_context=False - ) + quantum.GetQuantumJobRequest(name=JOB_PATH, return_run_context=False) ) - assert client.get_job('proj', 'prog', 'job0', True) == result + assert default_engine_client.get_job('proj', 'prog', 'job0', True) == result grpc_client.get_quantum_job.assert_called_with( - quantum.GetQuantumJobRequest( - name=JOB_PATH, return_run_context=True - ) + quantum.GetQuantumJobRequest(name=JOB_PATH, return_run_context=True) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_set_job_description(client_constructor): +def test_set_job_description(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumJob(name=JOB_PATH) grpc_client.update_quantum_job.return_value = result - client = EngineClient() - assert client.set_job_description('proj', 'prog', 'job0', 'A job') == result + assert default_engine_client.set_job_description('proj', 'prog', 'job0', 'A job') == result grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, - quantum_job=quantum.QuantumJob( - name=JOB_PATH, description='A job' - ), + quantum_job=quantum.QuantumJob(name=JOB_PATH, description='A job'), update_mask=FieldMask(paths=['description']), ) ) - assert client.set_job_description('proj', 'prog', 'job0', '') == result + assert default_engine_client.set_job_description('proj', 'prog', 'job0', '') == result grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, @@ -1070,7 +1078,7 @@ def test_set_job_description(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_set_job_labels(client_constructor): +def test_set_job_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) grpc_client.get_quantum_job.return_value = quantum.QuantumJob( @@ -1079,35 +1087,28 @@ def test_set_job_labels(client_constructor): result = quantum.QuantumJob(name=JOB_PATH) 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 default_engine_client.set_job_labels('proj', 'prog', 'job0', labels) == result grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, - quantum_job=quantum.QuantumJob( - name=JOB_PATH, - labels=labels, - label_fingerprint='hash', - ), + quantum_job=quantum.QuantumJob(name=JOB_PATH, labels=labels, label_fingerprint='hash'), update_mask=FieldMask(paths=['labels']), ) ) - assert client.set_job_labels('proj', 'prog', 'job0', {}) == result + assert default_engine_client.set_job_labels('proj', 'prog', 'job0', {}) == result grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, - quantum_job=quantum.QuantumJob( - name=JOB_PATH, label_fingerprint='hash' - ), + quantum_job=quantum.QuantumJob(name=JOB_PATH, label_fingerprint='hash'), update_mask=FieldMask(paths=['labels']), ) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_add_job_labels(client_constructor): +def test_add_job_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) existing = quantum.QuantumJob( @@ -1117,11 +1118,14 @@ def test_add_job_labels(client_constructor): result = quantum.QuantumJob(name=JOB_PATH) grpc_client.update_quantum_job.return_value = result - client = EngineClient() - assert client.add_job_labels('proj', 'prog', 'job0', {'color': 'red'}) == existing + assert ( + default_engine_client.add_job_labels('proj', 'prog', 'job0', {'color': 'red'}) == existing + ) assert grpc_client.update_quantum_job.call_count == 0 - assert client.add_job_labels('proj', 'prog', 'job0', {'hello': 'world'}) == result + assert ( + default_engine_client.add_job_labels('proj', 'prog', 'job0', {'hello': 'world'}) == result + ) grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, @@ -1135,7 +1139,10 @@ def test_add_job_labels(client_constructor): ) assert ( - client.add_job_labels('proj', 'prog', 'job0', {'hello': 'world', 'color': 'blue'}) == result + default_engine_client.add_job_labels( + 'proj', 'prog', 'job0', {'hello': 'world', 'color': 'blue'} + ) + == result ) grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( @@ -1151,7 +1158,7 @@ def test_add_job_labels(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_remove_job_labels(client_constructor): +def test_remove_job_labels(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) existing = quantum.QuantumJob( @@ -1161,73 +1168,71 @@ def test_remove_job_labels(client_constructor): result = quantum.QuantumJob(name=JOB_PATH) grpc_client.update_quantum_job.return_value = result - client = EngineClient() - assert client.remove_job_labels('proj', 'prog', 'job0', ['other']) == existing + assert default_engine_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 ( + default_engine_client.remove_job_labels('proj', 'prog', 'job0', ['hello', 'weather']) + == result + ) grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, quantum_job=quantum.QuantumJob( - name=JOB_PATH, - labels={'color': 'red', 'run': '1'}, - label_fingerprint='hash', + name=JOB_PATH, 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 + assert ( + default_engine_client.remove_job_labels('proj', 'prog', 'job0', ['color', 'weather', 'run']) + == result + ) grpc_client.update_quantum_job.assert_called_with( quantum.UpdateQuantumJobRequest( name=JOB_PATH, - quantum_job=quantum.QuantumJob( - name=JOB_PATH, label_fingerprint='hash' - ), + quantum_job=quantum.QuantumJob(name=JOB_PATH, label_fingerprint='hash'), update_mask=FieldMask(paths=['labels']), ) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_delete_job(client_constructor): +def test_delete_job(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) - client = EngineClient() - assert not client.delete_job('proj', 'prog', 'job0') + assert not default_engine_client.delete_job('proj', 'prog', 'job0') grpc_client.delete_quantum_job.assert_called_with( quantum.DeleteQuantumJobRequest(name=JOB_PATH) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_cancel_job(client_constructor): +def test_cancel_job(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) - client = EngineClient() - assert not client.cancel_job('proj', 'prog', 'job0') + assert not default_engine_client.cancel_job('proj', 'prog', 'job0') grpc_client.cancel_quantum_job.assert_called_with( quantum.CancelQuantumJobRequest(name=JOB_PATH) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_job_results(client_constructor): +def test_job_results(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumResult(parent=JOB_PATH) grpc_client.get_quantum_result.return_value = result - client = EngineClient() - assert client.get_job_results('proj', 'prog', 'job0') == result + assert default_engine_client.get_job_results('proj', 'prog', 'job0') == result grpc_client.get_quantum_result.assert_called_with( quantum.GetQuantumResultRequest(parent=JOB_PATH) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_jobs(client_constructor): +def test_list_jobs(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) results = [ @@ -1236,13 +1241,12 @@ def test_list_jobs(client_constructor): ] grpc_client.list_quantum_jobs.return_value = results - client = EngineClient() - assert client.list_jobs(project_id='proj', program_id='prog1') == results + assert default_engine_client.list_jobs(project_id='proj', program_id='prog1') == results grpc_client.list_quantum_jobs.assert_called_with( quantum.ListQuantumJobsRequest(parent='projects/proj/programs/prog1') ) - assert client.list_jobs(project_id='proj') == results + assert default_engine_client.list_jobs(project_id='proj') == results grpc_client.list_quantum_jobs.assert_called_with( quantum.ListQuantumJobsRequest(parent='projects/proj/programs/-') ) @@ -1340,10 +1344,10 @@ def test_list_jobs_filters( execution_states, executed_processor_ids, scheduled_processor_ids, + default_engine_client, ): grpc_client = _setup_client_mock(client_constructor) - client = EngineClient() - client.list_jobs( + default_engine_client.list_jobs( project_id='proj', program_id='prog', created_before=created_before, @@ -1369,7 +1373,7 @@ async def __aiter__(self): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_processors(client_constructor): +def test_list_processors(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) results = [ @@ -1378,29 +1382,27 @@ def test_list_processors(client_constructor): ] grpc_client.list_quantum_processors.return_value = Pager(results) - client = EngineClient() - assert client.list_processors('proj') == results + assert default_engine_client.list_processors('proj') == results grpc_client.list_quantum_processors.assert_called_with( quantum.ListQuantumProcessorsRequest(parent='projects/proj') ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_processor(client_constructor): +def test_get_processor(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumProcessor(name='projects/proj/processors/processor0') grpc_client.get_quantum_processor.return_value = result - client = EngineClient() - assert client.get_processor('proj', 'processor0') == result + assert default_engine_client.get_processor('proj', 'processor0') == result grpc_client.get_quantum_processor.assert_called_with( quantum.GetQuantumProcessorRequest(name='projects/proj/processors/processor0') ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_calibrations(client_constructor): +def test_list_calibrations(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) results = [ @@ -1409,15 +1411,14 @@ def test_list_calibrations(client_constructor): ] grpc_client.list_quantum_calibrations.return_value = Pager(results) - client = EngineClient() - assert client.list_calibrations('proj', 'processor0') == results + assert default_engine_client.list_calibrations('proj', 'processor0') == results grpc_client.list_quantum_calibrations.assert_called_with( quantum.ListQuantumCalibrationsRequest(parent='projects/proj/processors/processor0') ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_calibration(client_constructor): +def test_get_calibration(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumCalibration( @@ -1425,8 +1426,7 @@ def test_get_calibration(client_constructor): ) grpc_client.get_quantum_calibration.return_value = result - client = EngineClient() - assert client.get_calibration('proj', 'processor0', 123456) == result + assert default_engine_client.get_calibration('proj', 'processor0', 123456) == result grpc_client.get_quantum_calibration.assert_called_with( quantum.GetQuantumCalibrationRequest( name='projects/proj/processors/processor0/calibrations/123456' @@ -1435,7 +1435,7 @@ def test_get_calibration(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_current_calibration(client_constructor): +def test_get_current_calibration(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) result = quantum.QuantumCalibration( @@ -1443,8 +1443,7 @@ def test_get_current_calibration(client_constructor): ) grpc_client.get_quantum_calibration.return_value = result - client = EngineClient() - assert client.get_current_calibration('proj', 'processor0') == result + assert default_engine_client.get_current_calibration('proj', 'processor0') == result grpc_client.get_quantum_calibration.assert_called_with( quantum.GetQuantumCalibrationRequest( name='projects/proj/processors/processor0/calibrations/current' @@ -1453,13 +1452,12 @@ def test_get_current_calibration(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_current_calibration_does_not_exist(client_constructor): +def test_get_current_calibration_does_not_exist(client_constructor, default_engine_client): grpc_client = _setup_client_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 default_engine_client.get_current_calibration('proj', 'processor0') is None grpc_client.get_quantum_calibration.assert_called_with( quantum.GetQuantumCalibrationRequest( name='projects/proj/processors/processor0/calibrations/current' @@ -1468,24 +1466,22 @@ def test_get_current_calibration_does_not_exist(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_current_calibration_error(client_constructor): +def test_get_current_calibration_error(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) grpc_client.get_quantum_calibration.side_effect = exceptions.BadRequest('boom') - client = EngineClient() with pytest.raises(EngineException, match='boom'): - client.get_current_calibration('proj', 'processor0') + default_engine_client.get_current_calibration('proj', 'processor0') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_api_doesnt_retry_not_found_errors(client_constructor): +def test_api_doesnt_retry_not_found_errors(client_constructor, default_engine_client): grpc_client = _setup_client_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) + default_engine_client.get_program('proj', 'prog', False) assert grpc_client.get_quantum_program.call_count == 1 @@ -1516,7 +1512,7 @@ def test_api_retry_times(client_constructor, mock_sleep): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_create_reservation(client_constructor): +def test_create_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) start = datetime.datetime.fromtimestamp(1000000000) end = datetime.datetime.fromtimestamp(1000003600) @@ -1529,8 +1525,9 @@ def test_create_reservation(client_constructor): ) grpc_client.create_quantum_reservation.return_value = result - client = EngineClient() - assert client.create_reservation('proj', 'processor0', start, end, users) == result + assert ( + default_engine_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 = '' @@ -1542,7 +1539,7 @@ def test_create_reservation(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_cancel_reservation(client_constructor): +def test_cancel_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( @@ -1553,15 +1550,16 @@ def test_cancel_reservation(client_constructor): ) grpc_client.cancel_quantum_reservation.return_value = result - client = EngineClient() - assert client.cancel_reservation('proj', 'processor0', 'papar-party-44') == result + assert ( + default_engine_client.cancel_reservation('proj', 'processor0', 'papar-party-44') == result + ) grpc_client.cancel_quantum_reservation.assert_called_with( quantum.CancelQuantumReservationRequest(name=name) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_delete_reservation(client_constructor): +def test_delete_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( @@ -1572,15 +1570,16 @@ def test_delete_reservation(client_constructor): ) grpc_client.delete_quantum_reservation.return_value = result - client = EngineClient() - assert client.delete_reservation('proj', 'processor0', 'papar-party-44') == result + assert ( + default_engine_client.delete_reservation('proj', 'processor0', 'papar-party-44') == result + ) grpc_client.delete_quantum_reservation.assert_called_with( quantum.DeleteQuantumReservationRequest(name=name) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_reservation(client_constructor): +def test_get_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( @@ -1591,38 +1590,35 @@ def test_get_reservation(client_constructor): ) grpc_client.get_quantum_reservation.return_value = result - client = EngineClient() - assert client.get_reservation('proj', 'processor0', 'papar-party-44') == result + assert default_engine_client.get_reservation('proj', 'processor0', 'papar-party-44') == result grpc_client.get_quantum_reservation.assert_called_with( quantum.GetQuantumReservationRequest(name=name) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_reservation_not_found(client_constructor): +def test_get_reservation_not_found(client_constructor, default_engine_client): grpc_client = _setup_client_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') is None + assert default_engine_client.get_reservation('proj', 'processor0', 'papar-party-44') is None grpc_client.get_quantum_reservation.assert_called_with( quantum.GetQuantumReservationRequest(name=name) ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_get_reservation_exception(client_constructor): +def test_get_reservation_exception(client_constructor, default_engine_client): grpc_client = _setup_client_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') + default_engine_client.get_reservation('proj', 'processor0', 'goog') @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_reservation(client_constructor): +def test_list_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' results = [ @@ -1641,12 +1637,11 @@ def test_list_reservation(client_constructor): ] grpc_client.list_quantum_reservations.return_value = Pager(results) - client = EngineClient() - assert client.list_reservations('proj', 'processor0') == results + assert default_engine_client.list_reservations('proj', 'processor0') == results @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_update_reservation(client_constructor): +def test_update_reservation(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) name = 'projects/proj/processors/processor0/reservations/papar-party-44' result = quantum.QuantumReservation( @@ -1657,9 +1652,8 @@ def test_update_reservation(client_constructor): ) grpc_client.update_quantum_reservation.return_value = result - client = EngineClient() assert ( - client.update_reservation( + default_engine_client.update_reservation( 'proj', 'processor0', 'papar-party-44', @@ -1679,15 +1673,16 @@ def test_update_reservation(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_update_reservation_remove_all_users(client_constructor): +def test_update_reservation_remove_all_users(client_constructor, default_engine_client): grpc_client = _setup_client_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=[]) + default_engine_client.update_reservation( + 'proj', 'processor0', 'papar-party-44', whitelisted_users=[] + ) == result ) grpc_client.update_quantum_reservation.assert_called_with( @@ -1700,7 +1695,7 @@ def test_update_reservation_remove_all_users(client_constructor): @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) -def test_list_time_slots(client_constructor): +def test_list_time_slots(client_constructor, default_engine_client): grpc_client = _setup_client_mock(client_constructor) results = [ quantum.QuantumTimeSlot( @@ -1724,5 +1719,4 @@ def test_list_time_slots(client_constructor): ] grpc_client.list_quantum_time_slots.return_value = Pager(results) - client = EngineClient() - assert client.list_time_slots('proj', 'processor0') == results + assert default_engine_client.list_time_slots('proj', 'processor0') == results