From 7977b948702db7549878ebdca071cfe49a6ea2c6 Mon Sep 17 00:00:00 2001 From: Aaron Berdy Date: Wed, 4 Oct 2023 14:34:23 -0700 Subject: [PATCH] test: add job helpers to integ test (#724) This reverts commit 3d62becd9e297fabdedc66fe671fa86702728e61. --- test/integ_tests/job_test_script.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/integ_tests/job_test_script.py b/test/integ_tests/job_test_script.py index 7071ffd62..1fd7a364f 100644 --- a/test/integ_tests/job_test_script.py +++ b/test/integ_tests/job_test_script.py @@ -11,19 +11,19 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. -import json -import os - from braket.aws import AwsDevice from braket.circuits import Circuit -from braket.jobs import save_job_checkpoint, save_job_result +from braket.jobs import ( + get_hyperparameters, + get_job_device_arn, + save_job_checkpoint, + save_job_result, +) from braket.jobs_data import PersistedJobDataFormat def start_here(): - hp_file = os.environ["AMZN_BRAKET_HP_FILE"] - with open(hp_file, "r") as f: - hyperparameters = json.load(f) + hyperparameters = get_hyperparameters() if hyperparameters["test_case"] == "completed": completed_job_script() @@ -40,7 +40,7 @@ def completed_job_script(): print("Test job started!!!!!") # Use the device declared in the Orchestration Script - device = AwsDevice(os.environ["AMZN_BRAKET_DEVICE_ARN"]) + device = AwsDevice(get_job_device_arn()) bell = Circuit().h(0).cnot(0, 1) for count in range(5):