add mock for Quantum_Neuron

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

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/test/integ_tests/default_mocks/default_mocks.py b/test/integ_tests/default_mocks/default_mocks.py
@@ -1,10 +1,14 @@
 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()
-    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,10 +31,10 @@ 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):
-        self._wrapper.boto_client.get_default_jobs_role.return_value = mock.Mock()
+    def set_get_default_jobs_role(self, result):
+        self._wrapper.boto_client.get_default_jobs_role.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
@@ -3,8 +3,12 @@ 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.__iter__.return_value = []
+    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__))