From f6cc10687de9b9b0afa474701c8bc6663cd41b8a Mon Sep 17 00:00:00 2001
From: Ryan Shaffer <3620100+rmshaffer@users.noreply.github.com>
Date: Mon, 13 May 2024 17:19:39 -0400
Subject: [PATCH] Add tests for SerializableProgram handling

---
 .../braket/aws/test_aws_quantum_task.py       | 30 ++++++++++++++
 .../braket/devices/test_local_simulator.py    | 40 +++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/test/unit_tests/braket/aws/test_aws_quantum_task.py b/test/unit_tests/braket/aws/test_aws_quantum_task.py
index 16a72da7a..a63be2a96 100644
--- a/test/unit_tests/braket/aws/test_aws_quantum_task.py
+++ b/test/unit_tests/braket/aws/test_aws_quantum_task.py
@@ -33,6 +33,7 @@
     IRType,
     OpenQASMSerializationProperties,
     QubitReferenceType,
+    SerializableProgram,
 )
 from braket.device_schema import GateModelParameters, error_mitigation
 from braket.device_schema.dwave import (
@@ -123,6 +124,21 @@ def openqasm_program():
     return OpenQASMProgram(source="OPENQASM 3.0; h $0;")
 
 
+class DummySerializableProgram(SerializableProgram):
+    def __init__(self, source: str):
+        self.source = source
+
+    def to_ir(
+        self, ir_type: IRType = IRType.OPENQASM, allow_implicit_build: bool = False
+    ) -> str:
+        return self.source
+
+
+@pytest.fixture
+def serializable_program():
+    return DummySerializableProgram(source="OPENQASM 3.0; h $0;")
+
+
 @pytest.fixture
 def blackbird_program():
     return BlackbirdProgram(source="Vac | q[0]")
@@ -614,6 +630,20 @@ def test_create_openqasm_program_em_serialized(aws_session, arn, openqasm_progra
     )
 
 
+def test_create_serializable_program(aws_session, arn, serializable_program):
+    aws_session.create_quantum_task.return_value = arn
+    shots = 21
+    AwsQuantumTask.create(aws_session, SIMULATOR_ARN, serializable_program, S3_TARGET, shots)
+
+    _assert_create_quantum_task_called_with(
+        aws_session,
+        SIMULATOR_ARN,
+        OpenQASMProgram(source=serializable_program.to_ir()).json(),
+        S3_TARGET,
+        shots,
+    )
+
+
 def test_create_blackbird_program(aws_session, arn, blackbird_program):
     aws_session.create_quantum_task.return_value = arn
     shots = 21
diff --git a/test/unit_tests/braket/devices/test_local_simulator.py b/test/unit_tests/braket/devices/test_local_simulator.py
index a7e8bfe17..149f5555d 100644
--- a/test/unit_tests/braket/devices/test_local_simulator.py
+++ b/test/unit_tests/braket/devices/test_local_simulator.py
@@ -27,6 +27,7 @@
 from braket.annealing import Problem, ProblemType
 from braket.circuits import Circuit, FreeParameter, Gate, Noise
 from braket.circuits.noise_model import GateCriteria, NoiseModel, NoiseModelInstruction
+from braket.circuits.serialization import IRType, SerializableProgram
 from braket.device_schema import DeviceActionType, DeviceCapabilities
 from braket.device_schema.openqasm_device_action_properties import OpenQASMDeviceActionProperties
 from braket.devices import LocalSimulator, local_simulator
@@ -250,6 +251,26 @@ def properties(self) -> DeviceCapabilities:
         return device_properties
 
 
+class DummySerializableProgram(SerializableProgram):
+    def __init__(self, source: str):
+        self.source = source
+
+    def to_ir(
+        self, ir_type: IRType = IRType.OPENQASM, allow_implicit_build: bool = False
+    ) -> str:
+        return self.source
+
+
+class DummySerializableProgramSimulator(DummyProgramSimulator):
+    def run(
+        self,
+        program: SerializableProgram,
+        shots: int = 0,
+        batch_size: int = 1,
+    ) -> GateModelQuantumTaskResult:
+        return GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT)
+
+
 class DummyProgramDensityMatrixSimulator(BraketSimulator):
     def run(
         self, program: ir.openqasm.Program, shots: Optional[int], *args, **kwargs
@@ -556,6 +577,25 @@ def test_run_program_model():
     assert task.result() == GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT)
 
 
+def test_run_serializable_program_model():
+    dummy = DummySerializableProgramSimulator()
+    sim = LocalSimulator(dummy)
+    task = sim.run(
+        DummySerializableProgram(
+            source="""
+qubit[2] q;
+bit[2] c;
+
+h q[0];
+cnot q[0], q[1];
+
+c = measure q;
+"""
+        )
+    )
+    assert task.result() == GateModelQuantumTaskResult.from_object(GATE_MODEL_RESULT)
+
+
 @pytest.mark.xfail(raises=ValueError)
 def test_run_gate_model_value_error():
     dummy = DummyCircuitSimulator()