feat: mock getting started hybrid job example (#420)

Co-authored-by: Milan Krneta <krneta@amazon.com> Co-authored-by: Cody Wang <speller26@gmail.com>
amazon-braket · Oct 17, 2023 · 702f51e · 702f51e
1 parent 1ade6fc
commit 702f51e
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 39 deletions.
diff --git a/...b/Creating_your_first_Hybrid_Job_mocks.py → ...0_Creating_your_first_Hybrid_Job_mocks.py b/...b/Creating_your_first_Hybrid_Job_mocks.py → ...0_Creating_your_first_Hybrid_Job_mocks.py
@@ -1,8 +1,9 @@
 import os
+import sys
 import tarfile
 import subprocess
 import unittest.mock as mock
-
+from itertools import cycle
 
 default_job_results = ""
 
@@ -12,30 +13,57 @@ def pre_run_inject(mock_utils):
     mock_utils.mock_default_device_calls(mocker)
     mocker.set_search_result([
         {
-            "Roles" : [
+            "Roles": [
                 {
                     "RoleName": "AmazonBraketJobsExecutionRole",
-                    "Arn" : "TestRoleARN"
+                    "Arn": "TestRoleARN"
                 }
             ]
         }
     ])
-    mocker.set_create_job_result({
-        "jobArn" : f"arn:aws:braket:{mocker.region_name}:000000:job/testJob"
-    })
+    mocker.set_create_job_side_effect([
+        {
+            "jobArn": f"arn:aws:braket:{mocker.region_name}:000000:job/testJob"
+        },
+        {
+            "jobArn": f"arn:aws:braket:us-west-1:000000:job/testJob"
+        },
+    ])
     mocker.set_get_job_result({
-        "instanceConfig" : {
-            "instanceCount" : 1
+        "instanceConfig": {
+            "instanceCount": 1
         },
         "jobName": "testJob",
         "status": "COMPLETED",
         "outputDataConfig": {
-            "s3Path" : "s3://amazon-br-invalid-path/test-path/test-results"
+            "s3Path": "s3://amazon-br-invalid-path/test-path/test-results"
         }
     })
     mocker.set_log_streams_result({
         "logStreams": []
     })
+    mocker.set_get_query_results_result({
+        "status": "Complete",
+        "results": [
+            [
+                {"field": "@message", "value": "iteration_number=0;expval=0;"},
+                {"field": "@timestamp", "value": "0"},
+            ],
+        ]
+    })
+    mocker.set_batch_get_image_side_effect(
+        cycle([
+            {"images": [{"imageId": {"imageDigest": "my-digest"}}]},
+            {
+                "images": [
+                    {"imageId": {"imageTag": f"-py3{sys.version_info.minor}-"}},
+                ]
+            },
+        ])
+    )
+    mocker.set_start_query_result({
+        "queryId": "TestId"
+    })
     global default_job_results
     default_job_results = mock_utils.read_file("../job_results.json", __file__)
     with open("results.json", "w") as f:

diff --git a/test/integ_tests/mock_utils.py b/test/integ_tests/mock_utils.py
@@ -6,7 +6,6 @@
 import json
 import braket.aws
 
-
 plt.savefig = mock.Mock()
 
 
@@ -44,6 +43,9 @@ def set_search_result(self, result):
     def set_create_job_result(self, result):
         self._wrapper.boto_client.create_job.return_value = result
 
+    def set_create_job_side_effect(self, side_effect):
+        self._wrapper.boto_client.create_job.side_effect = side_effect
+
     def set_get_job_result(self, result):
         self._wrapper.boto_client.get_job.return_value = result
 
@@ -62,12 +64,15 @@ def set_get_query_results_result(self, result):
     def set_list_objects_v2_result(self, result):
         self._wrapper.boto_client.list_objects_v2.return_value = result
 
+    def set_batch_get_image_side_effect(self, side_effect):
+        self._wrapper.boto_client.batch_get_image.side_effect = side_effect
+
     @property
     def region_name(self):
         return self._wrapper.region_name
 
 
-def read_file(name, file_path = None):
+def read_file(name, file_path=None):
     if file_path:
         json_path = os.path.join(os.path.dirname(file_path), name)
     else:
@@ -78,33 +83,33 @@ def read_file(name, file_path = None):
 
 def mock_default_device_calls(mocker):
     mocker.set_get_device_result({
-        "deviceType" : "QPU",
-        "deviceCapabilities" : read_file("default_capabilities.json"),
+        "deviceType": "QPU",
+        "deviceCapabilities": read_file("default_capabilities.json"),
         "deviceQueueInfo": [
-        {
-        "queue": "QUANTUM_TASKS_QUEUE",
-        "queueSize": "13",
-        "queuePriority": "Normal"
-        },
-        {
-        "queue": "QUANTUM_TASKS_QUEUE",
-        "queueSize": "0",
-        "queuePriority": "Priority"
-        },
-        {
-        "queue": "JOBS_QUEUE",
-        "queueSize": "0"
-        }
-    ]
+            {
+                "queue": "QUANTUM_TASKS_QUEUE",
+                "queueSize": "13",
+                "queuePriority": "Normal"
+            },
+            {
+                "queue": "QUANTUM_TASKS_QUEUE",
+                "queueSize": "0",
+                "queuePriority": "Priority"
+            },
+            {
+                "queue": "JOBS_QUEUE",
+                "queueSize": "0"
+            }
+        ]
     })
     mocker.set_create_quantum_task_result({
-        "quantumTaskArn" : "arn:aws:braket:us-west-2:000000:quantum-task/TestARN",
+        "quantumTaskArn": "arn:aws:braket:us-west-2:000000:quantum-task/TestARN",
     })
     mocker.set_get_quantum_task_result({
-        "quantumTaskArn" : "arn:aws:braket:us-west-2:000000:quantum-task/TestARN",
-        "status" : "COMPLETED",
-        "outputS3Bucket" : "Test Bucket",
-        "outputS3Directory" : "Test Directory",
+        "quantumTaskArn": "arn:aws:braket:us-west-2:000000:quantum-task/TestARN",
+        "status": "COMPLETED",
+        "outputS3Bucket": "Test Bucket",
+        "outputS3Directory": "Test Directory",
         "shots": 10,
         "deviceArn": "Test Device Arn",
         "queueInfo": {
@@ -139,11 +144,10 @@ def __init__(self):
         self.boto_client.get_caller_identity.return_value = {
             "Account": "TestAccount"
         }
-        self.boto_client.meta.region_name = "us-west-2"
         self.boto_client.get_authorization_token.return_value = {
-            "authorizationData" : [
+            "authorizationData": [
                 {
-                    "authorizationToken" : "TestToken"
+                    "authorizationToken": "TestToken"
                 }
             ]
         }
@@ -153,8 +157,12 @@ class Boto3SessionAllWrapper(SessionWrapper):
     def __init__(self):
         super().__init__()
         boto3.Session = self
+        self._default_region = "us-west-2"
+        self._region = self._default_region
 
     def __call__(self, *args, **kwargs):
+        # handle explicit region_name=None
+        self._region = kwargs.get("region_name", None) or self._default_region
         return self
 
     def client(self, *args, **kwargs):
@@ -171,7 +179,7 @@ def get_credentials(self, *args, **kwargs):
 
     @property
     def region_name(self):
-        return "us-west-2"
+        return self._region
 
 
 class AwsSessionMinWrapper(SessionWrapper):
@@ -277,4 +285,3 @@ def retrieve_s3_object_body(self, s3_bucket, s3_object_key):
 
     def get_job_metrics(self, query_id):
         return AwsSessionFacade._wrapper.boto_client.get_query_results(query_id)["results"]
-
diff --git a/test/integ_tests/test_all_notebooks.py b/test/integ_tests/test_all_notebooks.py
@@ -22,7 +22,6 @@
 ]
 # Python 3.10 required for decorators notebooks
 EXCLUDED_NOTEBOOKS += [
-    "0_Creating_your_first_Hybrid_Job.ipynb",
     "Quantum_machine_learning_in_Amazon_Braket_Hybrid_Jobs.ipynb",
     "Using_PennyLane_with_Braket_Hybrid_Jobs.ipynb",
     "0_Getting_started.ipynb",