add mock for Quantum_Neuron

amazon-braket · Aug 12, 2023 · 3179173 · 3179173
1 parent a024e6a
commit 3179173
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/test/integ_tests/mock_utils.py b/test/integ_tests/mock_utils.py
@@ -27,6 +27,10 @@ def set_get_device_result(self, result):
 
     def set_create_quantum_task_result(self, result):
         self._wrapper.boto_client.create_quantum_task.return_value = result
+
+    ### test
+    def set_create_job_result(self, result):
+        self._wrapper.boto_client.create_job.return_value = result
 
     def set_get_quantum_task_result(self, result):
         self._wrapper.boto_client.get_quantum_task.return_value = result

diff --git a/test/integ_tests/textbook/Quantum_Neuron_mocks.py b/test/integ_tests/textbook/Quantum_Neuron_mocks.py
@@ -1,6 +1,11 @@
 def pre_run_inject(mock_utils):
     mocker = mock_utils.Mocker()
     mock_utils.mock_default_device_calls(mocker)
+    mocker.set_create_job_result(
+        {
+            "jobArn": "arn:aws:braket:us-east-1:000000:job/TestARN"
+        }
+    )   
     mocker.set_task_result_return(mock_utils.read_file("quantum_neuron_results.json", __file__))