From 514e4e062b16c678eda6b7b933674a4e36b7cb65 Mon Sep 17 00:00:00 2001
From: Serwan Asaad <serwan.asaad@quantum-machines.co>
Date: Wed, 3 Jul 2024 08:05:51 +0200
Subject: [PATCH] test ports containers

---
 quam/components/ports/ports_containers.py     | 11 ++-
 .../components/ports/test_ports_containers.py | 71 +++++++++++++++++++
 2 files changed, 76 insertions(+), 6 deletions(-)

diff --git a/quam/components/ports/ports_containers.py b/quam/components/ports/ports_containers.py
index 4805305a..bea29d54 100644
--- a/quam/components/ports/ports_containers.py
+++ b/quam/components/ports/ports_containers.py
@@ -83,18 +83,16 @@ def _get_port(
 
         return ports[port_id]
 
-    def get_port_from_reference(
-        self, port_reference: str, create=False
-    ) -> OPXPlusPortTypes:
+    def reference_to_port(self, port_reference: str, create=False) -> OPXPlusPortTypes:
         elems = port_reference.split("/")
         port_type, controller_id, port_id = elems[-3:]
 
+        port_type = port_type[:-1]
         if controller_id.isdigit():
             controller_id = int(controller_id)
-        fem_id = int(fem_id)
         port_id = int(port_id)
 
-        return self._get_port(controller_id, fem_id, port_id, port_type, create=create)
+        return self._get_port(controller_id, port_id, port_type, create=create)
 
     def get_analog_output(
         self,
@@ -217,7 +215,7 @@ def _get_port(
 
         return ports[port_id]
 
-    def get_port_from_reference(
+    def reference_to_port(
         self,
         port_reference: Union[QuamComponent, str],
         attr: Optional[str] = None,
@@ -232,6 +230,7 @@ def get_port_from_reference(
         elems = port_reference.split("/")
         port_type, controller_id, fem_id, port_id = elems[-4:]
 
+        port_type = port_type[:-1]
         if controller_id.isdigit():
             controller_id = int(controller_id)
         fem_id = int(fem_id)
diff --git a/tests/components/ports/test_ports_containers.py b/tests/components/ports/test_ports_containers.py
index bfbeb5d0..774e1ec0 100644
--- a/tests/components/ports/test_ports_containers.py
+++ b/tests/components/ports/test_ports_containers.py
@@ -69,6 +69,40 @@ def test_opx_plus_ports_container_add_ports(port_type):
     assert port is ports_group["con2"][3]
 
 
+@pytest.mark.parametrize(
+    "port_type", ["analog_output", "analog_input", "digital_output", "digital_input"]
+)
+def test_opx_plus_ports_container_reference_to_port(port_type):
+    port_mapping = {
+        "analog_output": OPXPlusAnalogOutputPort,
+        "analog_input": OPXPlusAnalogInputPort,
+        "digital_output": OPXPlusDigitalOutputPort,
+        "digital_input": OPXPlusDigitalInputPort,
+    }
+
+    ports_container = OPXPlusPortsContainer()
+
+    port_reference = f"#/{port_type}s/con1/3"
+
+    with pytest.raises(KeyError):
+        ports_container.reference_to_port(port_reference)
+
+    port = ports_container.reference_to_port(port_reference, create=True)
+    assert isinstance(port, port_mapping[port_type])
+
+    assert port.controller_id == "con1"
+    assert port.port_id == 3
+
+    port2 = ports_container.reference_to_port(port_reference, create=False)
+    assert port is port2
+
+    port3 = ports_container.reference_to_port(port_reference, create=True)
+    assert port is port3
+
+    ports_group = getattr(ports_container, f"{port_type}s")
+    assert port is ports_group["con1"][3]
+
+
 @pytest.mark.parametrize(
     "port_type",
     ["analog_output", "analog_input", "mw_output", "mw_input", "digital_output"],
@@ -104,3 +138,40 @@ def test_fem_ports_container_add_ports(port_type):
 
     ports_group = getattr(ports_container, f"{port_type}s")
     assert port is ports_group["con2"][2][3]
+
+
+@pytest.mark.parametrize(
+    "port_type",
+    ["analog_output", "analog_input", "mw_output", "mw_input", "digital_output"],
+)
+def test_fem_ports_container_reference_to_port(port_type):
+    port_mapping = {
+        "analog_output": LFFEMAnalogOutputPort,
+        "analog_input": LFFEMAnalogInputPort,
+        "mw_output": MWFEMAnalogOutputPort,
+        "mw_input": MWFEMAnalogInputPort,
+        "digital_output": FEMDigitalOutputPort,
+    }
+
+    port_reference = f"#/{port_type}s/con1/2/3"
+
+    ports_container = FEMPortsContainer()
+
+    with pytest.raises(KeyError):
+        ports_container.reference_to_port(port_reference)
+
+    port = ports_container.reference_to_port(port_reference, create=True)
+    assert isinstance(port, port_mapping[port_type])
+
+    assert port.controller_id == "con1"
+    assert port.fem_id == 2
+    assert port.port_id == 3
+
+    port2 = ports_container.reference_to_port(port_reference, create=False)
+    assert port is port2
+
+    port3 = ports_container.reference_to_port(port_reference, create=True)
+    assert port is port3
+
+    ports_group = getattr(ports_container, f"{port_type}s")
+    assert port is ports_group["con1"][2][3]