From d299db769a29bd68eac01eed2c288a4cea59c4b3 Mon Sep 17 00:00:00 2001
From: Dean Poulos <dean.poulos7@gmail.com>
Date: Fri, 11 Oct 2024 20:14:03 +1100
Subject: [PATCH] Write tests for channel reuse.:

---
 tests/wirer/test_block_channels.py      | 42 +++++++++++++++++++++++++
 tests/wirer/test_wirer_channel_reuse.py |  4 +--
 2 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 tests/wirer/test_block_channels.py

diff --git a/tests/wirer/test_block_channels.py b/tests/wirer/test_block_channels.py
new file mode 100644
index 00000000..607d4e99
--- /dev/null
+++ b/tests/wirer/test_block_channels.py
@@ -0,0 +1,42 @@
+from dataclasses import asdict
+
+import pytest
+
+from qualang_tools.wirer import *
+from qualang_tools.wirer.connectivity.element import QubitReference
+from qualang_tools.wirer.connectivity.wiring_spec import WiringLineType
+from qualang_tools.wirer.instruments.instrument_channel import InstrumentChannelLfFemOutput, \
+    InstrumentChannelMwFemInput, InstrumentChannelMwFemOutput
+
+visualize_flag = pytest.visualize_flag
+
+def test_alternating_blocking_of_used_channels(instruments_2lf_2mw):
+    connectivity = Connectivity()
+
+    connectivity.add_qubit_drive_lines(qubits=1)
+    allocate_wiring(connectivity, instruments_2lf_2mw, block_used_channels=False)
+
+    connectivity.add_qubit_drive_lines(qubits=2)
+    allocate_wiring(connectivity, instruments_2lf_2mw)
+
+    connectivity.add_qubit_drive_lines(qubits=3)
+    allocate_wiring(connectivity, instruments_2lf_2mw, block_used_channels=False)
+
+    connectivity.add_qubit_drive_lines(qubits=4)
+    allocate_wiring(connectivity, instruments_2lf_2mw)
+
+    if visualize_flag:
+        visualize(connectivity.elements, instruments_2lf_2mw.available_channels)
+
+    expected_ports = [
+        1,  # q1 allocated to 1, but channel isn't blocked
+        1,  # q2 allocated to 1, since it wasn't blocked
+        2,  # q3 allocated to 2, since it's the next available channel, but not blocked
+        2   # q4 allocated to 2, since it wasn't blocked
+    ]
+    for qubit_index in [1, 2, 3, 4]:
+        drive_channels = connectivity.elements[QubitReference(qubit_index)].channels[WiringLineType.DRIVE]
+        for i, channel in enumerate(drive_channels):
+            assert asdict(channel) == asdict([
+                InstrumentChannelMwFemOutput(con=1, slot=3, port=expected_ports[qubit_index-1])
+            ][i])
\ No newline at end of file
diff --git a/tests/wirer/test_wirer_channel_reuse.py b/tests/wirer/test_wirer_channel_reuse.py
index c3425651..3044ecf3 100644
--- a/tests/wirer/test_wirer_channel_reuse.py
+++ b/tests/wirer/test_wirer_channel_reuse.py
@@ -23,8 +23,8 @@ def test_5q_allocation_with_channel_reuse(instruments_2lf_2mw):
 
         allocate_wiring(connectivity, instruments_2lf_2mw, block_used_channels=False)
 
-    if visualize_flag:
-        visualize(connectivity.elements, instruments_2lf_2mw.available_channels)
+    # if visualize_flag:
+    visualize(connectivity.elements, instruments_2lf_2mw.available_channels)
 
     for qubit in [1, 3]:
         # resonator lines re-used for qubits 1 & 3