add mock for Quantum_Neuron

amazon-braket · Aug 12, 2023 · 853bd96 · 853bd96
1 parent 4a3ce2a
commit 853bd96
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 14 deletions.
diff --git a/test/integ_tests/default_mocks/default_mocks.py b/test/integ_tests/default_mocks/default_mocks.py
@@ -1,14 +1,10 @@
 def pre_run_inject(mock_utils):
     print("This is inside default_mocks.py")
-    
+
     mocker = mock_utils.Mocker()
     mock_utils.mock_default_device_calls(mocker)
-    mocker.set_get_default_jobs_role('arn:aws:iam::000000:role/service-role/TestRole')
-    mocker.set_create_job_result(
-        {
-            "jobArn": "arn:aws:braket:us-east-1:000000:job/TestARN"
-        }
-    )
+    mocker.set_get_default_jobs_role("arn:aws:iam::000000:role/service-role/TestRole")
+    mocker.set_create_job_result({"jobArn": "arn:aws:braket:us-east-1:000000:job/TestARN"})
 
 
 def post_run(tb):

diff --git a/test/integ_tests/mock_utils.py b/test/integ_tests/mock_utils.py
@@ -31,7 +31,7 @@ def set_create_quantum_task_result(self, result):
     # test
     def set_create_job_result(self, result):
         self._wrapper.boto_client.create_job.return_value = result
-    
+
     # test
     def set_get_default_jobs_role(self, result):
         self._wrapper.boto_client.get_default_jobs_role.return_value = result

diff --git a/test/integ_tests/textbook/Quantum_Neuron_mocks.py b/test/integ_tests/textbook/Quantum_Neuron_mocks.py
@@ -3,12 +3,8 @@ def pre_run_inject(mock_utils):
 
     mocker = mock_utils.Mocker()
     mock_utils.mock_default_device_calls(mocker)
-    mocker.set_get_default_jobs_role('arn:aws:iam::000000:role/service-role/TestRole')
-    mocker.set_create_job_result(
-        {
-            "jobArn": "arn:aws:braket:us-east-1:000000:job/TestARN"
-        }
-    )
+    mocker.set_get_default_jobs_role("arn:aws:iam::000000:role/service-role/TestRole")
+    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__))