From 702f51ec964a58e8493ae0d41da6692dcb0da18e Mon Sep 17 00:00:00 2001 From: Aaron Berdy Date: Tue, 17 Oct 2023 13:56:11 -0700 Subject: [PATCH] feat: mock getting started hybrid job example (#420) Co-authored-by: Milan Krneta Co-authored-by: Cody Wang --- ...0_Creating_your_first_Hybrid_Job_mocks.py} | 46 ++++++++++--- test/integ_tests/mock_utils.py | 65 ++++++++++--------- test/integ_tests/test_all_notebooks.py | 1 - 3 files changed, 73 insertions(+), 39 deletions(-) rename test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/{Creating_your_first_Hybrid_Job_mocks.py => 0_Creating_your_first_Hybrid_Job_mocks.py} (63%) diff --git a/test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/Creating_your_first_Hybrid_Job_mocks.py b/test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/0_Creating_your_first_Hybrid_Job_mocks.py similarity index 63% rename from test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/Creating_your_first_Hybrid_Job_mocks.py rename to test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/0_Creating_your_first_Hybrid_Job_mocks.py index 855737011..4377671e2 100644 --- a/test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/Creating_your_first_Hybrid_Job_mocks.py +++ b/test/integ_tests/hybrid_jobs/0_Creating_your_first_Hybrid_Job/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 index 625efafae..58790ad1a 100644 --- 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 index 3fc829e45..667768cc0 100644 --- 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",