GetConnection(s) to Support Pin Groups (#639)

* feat: support get connection(s) for pins group(s)

Author: Avinash2Suresh

ni_measurementlink_service/session_management/_client.py

@@ -5,7 +5,7 @@
 import logging
 import threading
 import warnings
-from typing import Dict, Iterable, Optional, Union
+from typing import Dict, Iterable, Mapping, Optional, Union
 
 import google.protobuf.internal.containers
 import grpc
@@ -398,7 +398,7 @@ def _to_group_mappings_dict(
     mappings: google.protobuf.internal.containers.MessageMap[
         str, session_management_service_pb2.ResolvedPinsOrRelays
     ]
-) -> Dict[str, Iterable[str]]:
+) -> Mapping[str, Iterable[str]]:
     group_mappings: Dict[str, Iterable[str]] = {}
     if mappings is not None:
         for key, value in mappings.items():

ni_measurementlink_service/session_management/_reservation.py

@@ -20,6 +20,7 @@
     Iterable,
     List,
     Literal,
+    Mapping,
     NamedTuple,
     Optional,
     Sequence,
@@ -568,7 +569,7 @@ def __init__(
         multiplexer_session_info: Optional[
             Sequence[session_management_service_pb2.MultiplexerSessionInformation]
         ] = None,
-        pin_or_relay_group_mappings: Optional[Dict[str, Iterable[str]]] = None,
+        pin_or_relay_group_mappings: Optional[Mapping[str, Iterable[str]]] = None,
         reserved_pin_or_relay_names: Union[str, Iterable[str], None] = None,
         reserved_sites: Optional[Iterable[int]] = None,
     ) -> None:
@@ -583,7 +584,7 @@ def __init__(
         self._multiplexer_session_container = MultiplexerSessionContainer(
             session_management_client, multiplexer_session_info
         )
-        self._pin_or_relay_group_mappings: Dict[str, Iterable[str]] = {}
+        self._pin_or_relay_group_mappings: Mapping[str, Iterable[str]] = {}
         if pin_or_relay_group_mappings is not None:
             self._pin_or_relay_group_mappings = pin_or_relay_group_mappings
 
@@ -810,19 +811,27 @@ def _get_connections_core(
     ) -> Sequence[TypedConnection[TSession]]:
         _check_optional_str_param("instrument_type_id", instrument_type_id)
 
-        requested_pins = _to_iterable(pin_or_relay_names, self._reserved_pin_or_relay_names)
+        requested_pin_or_relay_names = _to_iterable(
+            pin_or_relay_names, self._reserved_pin_or_relay_names
+        )
         requested_sites = _to_iterable(sites, self._reserved_sites)
         requested_instrument_type_ids = _to_iterable(
             instrument_type_id, self._reserved_instrument_type_ids
         )
 
+        resolved_pin_or_relay_names = _to_ordered_set(
+            self._get_resolved_pin_or_relay_names(requested_pin_or_relay_names)
+        )
+
         # Validate that each requested pin, site, or instrument type ID is
         # present in the reserved pins, reserved sites, and reserved instrument
         # type IDs. This rejects unknown or invalid inputs such as
         # pin_or_relay_names="NonExistentPin" or sites=[0, 1, 65535].
         if pin_or_relay_names is not None:
             _check_matching_criterion(
-                "pin or relay name(s)", requested_pins, self._reserved_pin_or_relay_names
+                "pin or relay name(s)",
+                resolved_pin_or_relay_names,
+                self._reserved_pin_or_relay_names,
             )
         if sites is not None:
             _check_matching_criterion("site(s)", requested_sites, self._reserved_sites)
@@ -842,7 +851,7 @@ def _get_connections_core(
         results: List[TypedConnection[TSession]] = []
         matching_pins: Set[str] = set()
         for site in requested_sites_with_system:
-            for pin in requested_pins:
+            for pin in resolved_pin_or_relay_names:
                 for instrument_type in requested_instrument_type_ids:
                     key = _ConnectionKey(pin, site, instrument_type)
                     value = self._connection_cache.get(key)
@@ -862,10 +871,10 @@ def _get_connections_core(
 
         # If the user specified pins to match, validate that each one matched a connection.
         if pin_or_relay_names is not None and not all(
-            pin in matching_pins for pin in requested_pins
+            pin in matching_pins for pin in resolved_pin_or_relay_names
         ):
             extra_pins_str = ", ".join(
-                _quote(pin) for pin in requested_pins if pin not in matching_pins
+                _quote(pin) for pin in resolved_pin_or_relay_names if pin not in matching_pins
             )
             criteria = _describe_matching_criteria(None, sites, instrument_type_id)
             # Emphasize the extra pin/relay names, but also list the other criteria.

tests/integration/session_management/test_reservation.py

@@ -254,6 +254,92 @@ def test___sessions_reserved_using_nested_relay_group___get_connections_by_relay
         ]
 
 
+def test___sessions_reserved___get_connections_by_pin_group___returns_connections(
+    pin_map_client: PinMapClient,
+    pin_map_directory: pathlib.Path,
+    session_management_client: SessionManagementClient,
+) -> None:
+    with ExitStack() as stack:
+        pin_map_id = pin_map_client.update_pin_map(pin_map_directory / _PIN_MAP_C)
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=[0])
+        pin_group = ["PinGroup1"]
+        reservation = stack.enter_context(
+            session_management_client.reserve_sessions(pin_map_context, pin_group)
+        )
+
+        connections = reservation.get_connections(object, pin_or_relay_names=pin_group)
+
+        assert [get_connection_subset(conn) for conn in connections] == [
+            ConnectionSubset("A", 0, "DCPower1/0, DCPower1/2, DCPower2/1", "DCPower1/0"),
+            ConnectionSubset("S1", -1, "SCOPE1", "1"),
+        ]
+
+
+def test___sessions_reserved___get_connections_by_nested_pin_group___returns_connections(
+    pin_map_client: PinMapClient,
+    pin_map_directory: pathlib.Path,
+    session_management_client: SessionManagementClient,
+) -> None:
+    with ExitStack() as stack:
+        pin_map_id = pin_map_client.update_pin_map(pin_map_directory / _PIN_MAP_C)
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=[0])
+        pin_groups = ["PinGroup1", "PinGroup2"]
+        reservation = stack.enter_context(
+            session_management_client.reserve_sessions(pin_map_context, pin_groups)
+        )
+
+        connections = reservation.get_connections(object, pin_or_relay_names=pin_groups)
+
+        assert [get_connection_subset(conn) for conn in connections] == [
+            ConnectionSubset("A", 0, "DCPower1/0, DCPower1/2, DCPower2/1", "DCPower1/0"),
+            ConnectionSubset("C", 0, "SCOPE1", "2"),
+            ConnectionSubset("S1", -1, "SCOPE1", "1"),
+        ]
+
+
+def test___sessions_reserved___get_connections_by_relay_group___returns_connections(
+    pin_map_client: PinMapClient,
+    pin_map_directory: pathlib.Path,
+    session_management_client: SessionManagementClient,
+) -> None:
+    with ExitStack() as stack:
+        pin_map_id = pin_map_client.update_pin_map(pin_map_directory / _PIN_MAP_C)
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=[0])
+        relay_group = ["RelayGroup1"]
+        reservation = stack.enter_context(
+            session_management_client.reserve_sessions(pin_map_context, relay_group)
+        )
+
+        connections = reservation.get_connections(object, pin_or_relay_names=relay_group)
+
+        assert [get_connection_subset(conn) for conn in connections] == [
+            ConnectionSubset("RelayUsingSameDriver", 0, "RelayDriver1", "K0"),
+            ConnectionSubset("SystemRelay", -1, "RelayDriver1", "K60"),
+        ]
+
+
+def test___sessions_reserved___get_connections_by_nested_relay_groups___returns_connections(
+    pin_map_client: PinMapClient,
+    pin_map_directory: pathlib.Path,
+    session_management_client: SessionManagementClient,
+) -> None:
+    with ExitStack() as stack:
+        pin_map_id = pin_map_client.update_pin_map(pin_map_directory / _PIN_MAP_C)
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=[0])
+        relay_groups = ["RelayGroup1", "RelayGroup2"]
+        reservation = stack.enter_context(
+            session_management_client.reserve_sessions(pin_map_context, relay_groups)
+        )
+
+        connections = reservation.get_connections(object, pin_or_relay_names=relay_groups)
+
+        assert [get_connection_subset(conn) for conn in connections] == [
+            ConnectionSubset("RelayUsingSameDriver", 0, "RelayDriver1", "K0"),
+            ConnectionSubset("RelayUsingDifferentDrivers", 0, "RelayDriver1", "K10"),
+            ConnectionSubset("SystemRelay", -1, "RelayDriver1", "K60"),
+        ]
+
+
 def test___reserve_sessions_with_multiplexer___get_connections_with_multiplexer___returns_connections(
     pin_map_client: PinMapClient,
     pin_map_directory: pathlib.Path,

tests/unit/test_reservation.py

@@ -1,6 +1,6 @@
 import functools
 from contextlib import ExitStack
-from typing import Any, Dict, Iterable, List
+from typing import Any, Dict, List
 from unittest.mock import Mock
 
 import pytest
@@ -242,7 +242,7 @@ def test___session_reserved_using_pin_group___get_connection_by_pin___returns_co
         grpc_session_infos = create_nifake_session_infos(1)
         grpc_session_infos[0].channel_mappings.add(pin_or_relay_name="Pin1", site=2, channel="3")
         grpc_session_infos[0].channel_mappings.add(pin_or_relay_name="Pin2", site=2, channel="2")
-        group_mappings: Dict[str, Iterable[str]] = {"PinGroup1": ["Pin1", "Pin2"]}
+        group_mappings = {"PinGroup1": ["Pin1", "Pin2"]}
         reservation = MultiSessionReservation(
             session_management_client,
             grpc_session_infos,
@@ -261,6 +261,54 @@ def test___session_reserved_using_pin_group___get_connection_by_pin___returns_co
         assert connection.session_info == session_info
 
 
+def test___duplicate_pins___get_connections___returns_single_connection(
+    session_management_client: Mock,
+) -> None:
+    with ExitStack() as stack:
+        grpc_session_infos = create_nifake_session_infos(1)
+        grpc_session_infos[0].channel_mappings.add(pin_or_relay_name="Pin1", site=2, channel="3")
+        reservation = MultiSessionReservation(session_management_client, grpc_session_infos)
+        session_info = stack.enter_context(
+            reservation.initialize_session(construct_session, "nifake")
+        )
+
+        connections = reservation.get_connections(
+            fake_driver.Session, pin_or_relay_names=["Pin1", "Pin1"]
+        )
+
+        assert len(connections) == 1
+        assert connections[0].pin_or_relay_name == "Pin1"
+        assert connections[0].site == 2
+        assert connections[0].channel_name == "3"
+        assert connections[0].session_info == session_info
+
+
+def test___session_reserved___get_connections_by_pin_group___returns_connections(
+    session_management_client: Mock,
+) -> None:
+    with ExitStack() as stack:
+        grpc_session_infos = create_nifake_session_infos(1)
+        grpc_session_infos[0].channel_mappings.add(pin_or_relay_name="Pin1", site=2, channel="3")
+        grpc_session_infos[0].channel_mappings.add(pin_or_relay_name="Pin2", site=2, channel="2")
+        group_mappings = {"PinGroup1": ["Pin1", "Pin2"]}
+        reservation = MultiSessionReservation(
+            session_management_client,
+            grpc_session_infos,
+            pin_or_relay_group_mappings=group_mappings,
+            reserved_pin_or_relay_names=["PinGroup1"],
+        )
+        _ = stack.enter_context(reservation.initialize_session(construct_session, "nifake"))
+
+        connections = reservation.get_connections(
+            fake_driver.Session, pin_or_relay_names=["PinGroup1"]
+        )
+
+        assert [get_connection_subset(conn) for conn in connections] == [
+            ConnectionSubset("Pin1", 2, "Dev0", "3"),
+            ConnectionSubset("Pin2", 2, "Dev0", "2"),
+        ]
+
+
 def test___multiple_connections___get_connection___value_error_raised(
     session_management_client: Mock,
 ) -> None: