use mock DM1 in unit tests

amazon-braket · Jan 26, 2024 · 86cdb38 · 86cdb38
1 parent c50d8ac
commit 86cdb38
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/test/unit_tests/braket/aws/common_test_utils.py b/test/unit_tests/braket/aws/common_test_utils.py
@@ -21,6 +21,7 @@
 IONQ_ARN = "arn:aws:braket:us-east-1::device/qpu/ionq/Harmony"
 OQC_ARN = "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy"
 SV1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
+DM1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/dm1"
 TN1_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/tn1"
 XANADU_ARN = "arn:aws:braket:us-east-1::device/qpu/xanadu/Borealis"
 

diff --git a/test/unit_tests/braket/aws/test_aws_device.py b/test/unit_tests/braket/aws/test_aws_device.py
@@ -22,6 +22,7 @@
 import pytest
 from botocore.exceptions import ClientError
 from common_test_utils import (
+    DM1_ARN,
     DWAVE_ARN,
     IONQ_ARN,
     OQC_ARN,
@@ -362,6 +363,27 @@ def test_d_wave_schema():
             "version": ["1"],
             "supportedOperations": ["H"],
         },
+    },
+    "paradigm": {"qubitCount": 30},
+    "deviceParameters": {},
+}
+
+MOCK_GATE_MODEL_NOISE_SIMULATOR_CAPABILITIES_JSON = {
+    "braketSchemaHeader": {
+        "name": "braket.device_schema.simulators.gate_model_simulator_device_capabilities",
+        "version": "1",
+    },
+    "service": {
+        "executionWindows": [
+            {
+                "executionDay": "Everyday",
+                "windowStartHour": "11:00",
+                "windowEndHour": "12:00",
+            }
+        ],
+        "shotsRange": [1, 10],
+    },
+    "action": {
         "braket.ir.openqasm.program": {
             "actionType": "braket.ir.openqasm.program",
             "version": ["1"],
@@ -409,6 +431,18 @@ def test_gate_model_sim_schema():
     )
 
 
+MOCK_GATE_MODEL_NOISE_SIMULATOR_CAPABILITIES = GateModelSimulatorDeviceCapabilities.parse_obj(
+    MOCK_GATE_MODEL_NOISE_SIMULATOR_CAPABILITIES_JSON
+)
+
+
+def test_gate_model_sim_schema():
+    validate(
+        MOCK_GATE_MODEL_NOISE_SIMULATOR_CAPABILITIES_JSON,
+        GateModelSimulatorDeviceCapabilities.schema(),
+    )
+
+
 MOCK_GATE_MODEL_SIMULATOR = {
     "deviceName": "SV1",
     "deviceType": "SIMULATOR",
@@ -417,6 +451,16 @@ def test_gate_model_sim_schema():
     "deviceCapabilities": MOCK_GATE_MODEL_SIMULATOR_CAPABILITIES.json(),
 }
 
+
+MOCK_GATE_MODEL_NOISE_SIMULATOR = {
+    "deviceName": "DM1",
+    "deviceType": "SIMULATOR",
+    "providerName": "provider1",
+    "deviceStatus": "ONLINE",
+    "deviceCapabilities": MOCK_GATE_MODEL_NOISE_SIMULATOR_CAPABILITIES.json(),
+}
+
+
 MOCK_DEFAULT_S3_DESTINATION_FOLDER = (
     "amazon-braket-us-test-1-00000000",
     "tasks",
@@ -2121,9 +2165,9 @@ def noise_model():
 @patch("braket.aws.aws_device.AwsSession")
 @patch("braket.aws.aws_quantum_task.AwsQuantumTask.create")
 def test_run_with_noise_model(aws_quantum_task_mock, aws_session_init, aws_session, noise_model):
-    arn = SV1_ARN
+    arn = DM1_ARN
     aws_session_init.return_value = aws_session
-    aws_session.get_device.return_value = MOCK_GATE_MODEL_SIMULATOR
+    aws_session.get_device.return_value = MOCK_GATE_MODEL_NOISE_SIMULATOR
     device = AwsDevice(arn, noise_model=noise_model)
     circuit = Circuit().h(0).cnot(0, 1)
     _ = device.run(circuit)