From 077a42ea0aeb127ae746098e82d46a339cf2af5f Mon Sep 17 00:00:00 2001
From: Phillip Weinberg <weinbe58@gmail.com>
Date: Fri, 29 Sep 2023 08:42:50 -0400
Subject: [PATCH] adding unit test.

---
 src/bloqade/factory.py        |  2 +
 tests/test_builder_factory.py | 98 +++++++++++++++++++++++++++++++++--
 2 files changed, 96 insertions(+), 4 deletions(-)

diff --git a/src/bloqade/factory.py b/src/bloqade/factory.py
index c360d79b8..05b236b78 100644
--- a/src/bloqade/factory.py
+++ b/src/bloqade/factory.py
@@ -54,6 +54,8 @@ def rydberg_h(
     from bloqade import start
     from bloqade.atom_arrangement import AtomArrangement
 
+    print(type(atoms_positions))
+
     if isinstance(atoms_positions, AtomArrangement):
         prog = atoms_positions
     else:
diff --git a/tests/test_builder_factory.py b/tests/test_builder_factory.py
index ae24fad47..65eb3728c 100644
--- a/tests/test_builder_factory.py
+++ b/tests/test_builder_factory.py
@@ -1,11 +1,28 @@
-"""
-from bloqade.builder.factory import (
+from bloqade import (
+    waveform,
+    rydberg_h,
     piecewise_linear,
     piecewise_constant,
     constant,
     linear,
+    var,
+    cast,
+    start,
 )
-from bloqade import cast
+from bloqade.atom_arrangement import Chain
+from bloqade.ir import (
+    AnalogCircuit,
+    Sequence,
+    rydberg,
+    Pulse,
+    rabi,
+    detuning,
+    Field,
+    Uniform,
+)
+from bloqade.ir.routine.params import Params
+from bloqade.ir.routine.base import Routine
+import numpy as np
 
 
 def test_ir_piecewise_linear():
@@ -54,4 +71,77 @@ def test_ir_piecewise_constant():
 
     assert A.waveforms[2].duration == cast(0.2)
     assert A.waveforms[2].value == cast("b")
-"""
+
+
+def test_rydberg_h():
+    run_time = var("run_time")
+
+    @waveform(run_time + 0.2)
+    def delta(t, amp, omega):
+        return np.sin(omega * t) * amp
+
+    delta = delta.sample(0.05, "linear")
+    ampl = piecewise_linear([0.1, run_time, 0.1], [0, 10, 10, 0])
+    phase = piecewise_constant([2, 2], [0, np.pi])
+    register = Chain(11, lattice_spacing=6.1)
+
+    static_params = {"amp": 1.0}
+    batch_params = [{"omega": omega} for omega in [1, 2, 4, 8]]
+    args = ["run_time"]
+
+    prog = rydberg_h(
+        register,
+        detuning=delta,
+        amplitude=ampl,
+        phase=phase,
+        batch_params=batch_params,
+        static_params=static_params,
+        args=args,
+    )
+
+    detuning_field = Field({Uniform: delta})
+    ampl_field = Field({Uniform: ampl})
+    phase_field = Field({Uniform: phase})
+
+    pulse = Pulse(
+        {detuning: detuning_field, rabi.amplitude: ampl_field, rabi.phase: phase_field}
+    )
+    sequence = Sequence({rydberg: pulse})
+
+    routine = prog.parse()
+
+    circuit = AnalogCircuit(register, sequence)
+    params = Params(static_params, batch_params, args)
+
+    assert routine == Routine(prog, circuit, params)
+
+
+def test_rydberg_h_2():
+    run_time = var("run_time")
+
+    @waveform(run_time + 0.2)
+    def delta(t, amp, omega):
+        return np.sin(omega * t) * amp
+
+    delta = delta.sample(0.05, "linear")
+    ampl = piecewise_linear([0.1, run_time, 0.1], [0, 10, 10, 0])
+    phase = piecewise_constant([2, 2], [0, np.pi])
+    register = start.add_position((0, 0))
+
+    prog = rydberg_h(
+        (0, 0), detuning=delta, amplitude=ampl, phase=phase, batch_params={}
+    )
+
+    detuning_field = Field({Uniform: delta})
+    ampl_field = Field({Uniform: ampl})
+    phase_field = Field({Uniform: phase})
+
+    pulse = Pulse(
+        {detuning: detuning_field, rabi.amplitude: ampl_field, rabi.phase: phase_field}
+    )
+    sequence = Sequence({rydberg: pulse})
+
+    print(prog.parse_circuit())
+    print(AnalogCircuit(register, sequence))
+
+    assert prog.parse_circuit() == AnalogCircuit(register, sequence)