examples: Update TestStand code modules for NI-DAQmx, NI-DCPower and NI-DMM examples

examples/nidaqmx_analog_input/teststand_nidaqmx.py

@@ -1,10 +1,14 @@
 """Functions to set up and tear down sessions of NI-DAQmx devices in NI TestStand."""
 from typing import Any
 
-import ni_measurementlink_service as nims
-import nidaqmx
 from _helpers import GrpcChannelPoolHelper, TestStandSupport
-from _nidaqmx_helpers import create_task
+from ni_measurementlink_service.discovery import DiscoveryClient
+from ni_measurementlink_service.session_management import (
+    INSTRUMENT_TYPE_NI_DAQMX,
+    PinMapContext,
+    SessionInitializationBehavior,
+    SessionManagementClient,
+)
 
 
 def create_nidaqmx_tasks(sequence_context: Any) -> None:
@@ -15,49 +19,42 @@ def create_nidaqmx_tasks(sequence_context: Any) -> None:
             (Dynamically typed.)
     """
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
         teststand_support = TestStandSupport(sequence_context)
         pin_map_id = teststand_support.get_active_pin_map_id()
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=None)
 
-        pin_map_context = nims.session_management.PinMapContext(pin_map_id=pin_map_id, sites=None)
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidaqmx.GRPC_SERVICE_INTERFACE_NAME
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
+
         with session_management_client.reserve_sessions(
-            context=pin_map_context,
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DAQMX,
+            pin_map_context, instrument_type_id=INSTRUMENT_TYPE_NI_DAQMX
         ) as reservation:
-            for session_info in reservation.session_info:
-                task = create_task(
-                    session_info,
-                    grpc_device_channel,
-                    initialization_behavior=nidaqmx.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
-                )
-                task.ai_channels.add_ai_voltage_chan(session_info.channel_list)
+            with reservation.create_nidaqmx_tasks(
+                initialization_behavior=SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH
+            ) as tasks:
+                for task in tasks:
+                    task.session.ai_channels.add_ai_voltage_chan(task.channel_list)
 
             session_management_client.register_sessions(reservation.session_info)
 
 
 def destroy_nidaqmx_tasks() -> None:
     """Destroy and unregister all NI-DAQmx tasks."""
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidaqmx.GRPC_SERVICE_INTERFACE_NAME
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
         with session_management_client.reserve_all_registered_sessions(
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DAQMX,
+            instrument_type_id=INSTRUMENT_TYPE_NI_DAQMX,
         ) as reservation:
-            session_management_client.unregister_sessions(reservation.session_info)
+            if not reservation.session_info:
+                return
 
-            for session_info in reservation.session_info:
-                task = create_task(
-                    session_info,
-                    grpc_device_channel,
-                    initialization_behavior=nidaqmx.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
-                )
-                task.close()
+            session_management_client.unregister_sessions(reservation.session_info)
+            with reservation.create_nidaqmx_tasks(
+                initialization_behavior=SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE
+            ):
+                pass

examples/nidcpower_source_dc_voltage/teststand_nidcpower.py

@@ -1,11 +1,14 @@
 """Functions to set up and tear down sessions of NI-DCPower devices in NI TestStand."""
 from typing import Any
 
-import ni_measurementlink_service as nims
-import nidcpower
-from _constants import USE_SIMULATION
 from _helpers import GrpcChannelPoolHelper, TestStandSupport
-from _nidcpower_helpers import create_session
+from ni_measurementlink_service.discovery import DiscoveryClient
+from ni_measurementlink_service.session_management import (
+    INSTRUMENT_TYPE_NI_DCPOWER,
+    PinMapContext,
+    SessionInitializationBehavior,
+    SessionManagementClient,
+)
 
 
 def create_nidcpower_sessions(sequence_context: Any) -> None:
@@ -16,50 +19,40 @@ def create_nidcpower_sessions(sequence_context: Any) -> None:
             (Dynamically typed.)
     """
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
-
         teststand_support = TestStandSupport(sequence_context)
         pin_map_id = teststand_support.get_active_pin_map_id()
-        pin_map_context = nims.session_management.PinMapContext(pin_map_id=pin_map_id, sites=None)
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidcpower.GRPC_SERVICE_INTERFACE_NAME
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=None)
+
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
         with session_management_client.reserve_sessions(
-            context=pin_map_context,
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DCPOWER,
+            pin_map_context, instrument_type_id=INSTRUMENT_TYPE_NI_DCPOWER
         ) as reservation:
-            for session_info in reservation.session_info:
-                # Leave session open
-                _ = create_session(
-                    session_info,
-                    USE_SIMULATION,
-                    grpc_device_channel,
-                    initialization_behavior=nidcpower.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
-                )
+            with reservation.initialize_nidcpower_sessions(
+                initialization_behavior=SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH
+            ):
+                pass
+
             session_management_client.register_sessions(reservation.session_info)
 
 
 def destroy_nidcpower_sessions() -> None:
     """Destroy and unregister all NI-DCPower sessions."""
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidcpower.GRPC_SERVICE_INTERFACE_NAME
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
         with session_management_client.reserve_all_registered_sessions(
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DCPOWER,
+            instrument_type_id=INSTRUMENT_TYPE_NI_DCPOWER,
         ) as reservation:
-            session_management_client.unregister_sessions(reservation.session_info)
+            if not reservation.session_info:
+                return
 
-            for session_info in reservation.session_info:
-                session = create_session(
-                    session_info,
-                    USE_SIMULATION,
-                    grpc_device_channel,
-                    initialization_behavior=nidcpower.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
-                )
-                session.close()
+            session_management_client.unregister_sessions(reservation.session_info)
+            with reservation.initialize_nidcpower_sessions(
+                initialization_behavior=SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE
+            ):
+                pass

examples/nidmm_measurement/teststand_nidmm.py

@@ -1,10 +1,14 @@
 """Functions to set up and tear down sessions of NI-DMM devices in NI TestStand."""
 from typing import Any
 
-import ni_measurementlink_service as nims
-import nidmm
 from _helpers import GrpcChannelPoolHelper, TestStandSupport
-from _nidmm_helpers import create_session
+from ni_measurementlink_service.discovery import DiscoveryClient
+from ni_measurementlink_service.session_management import (
+    INSTRUMENT_TYPE_NI_DMM,
+    PinMapContext,
+    SessionInitializationBehavior,
+    SessionManagementClient,
+)
 
 
 def create_nidmm_sessions(sequence_context: Any) -> None:
@@ -15,48 +19,41 @@ def create_nidmm_sessions(sequence_context: Any) -> None:
             (Dynamically typed.)
     """
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
-
         teststand_support = TestStandSupport(sequence_context)
         pin_map_id = teststand_support.get_active_pin_map_id()
-        pin_map_context = nims.session_management.PinMapContext(pin_map_id=pin_map_id, sites=None)
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidmm.GRPC_SERVICE_INTERFACE_NAME
+        pin_map_context = PinMapContext(pin_map_id=pin_map_id, sites=None)
+
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
+
         with session_management_client.reserve_sessions(
-            context=pin_map_context,
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DMM,
+            pin_map_context, instrument_type_id=INSTRUMENT_TYPE_NI_DMM
         ) as reservation:
-            for session_info in reservation.session_info:
-                # Leave session open
-                _ = create_session(
-                    session_info,
-                    grpc_device_channel,
-                    initialization_behavior=nidmm.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
-                )
+            with reservation.initialize_nidmm_sessions(
+                initialization_behavior=SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH
+            ):
+                pass
 
             session_management_client.register_sessions(reservation.session_info)
 
 
 def destroy_nidmm_sessions() -> None:
     """Destroy and unregister all NI-DMM sessions."""
     with GrpcChannelPoolHelper() as grpc_channel_pool:
-        session_management_client = nims.session_management.Client(
-            grpc_channel=grpc_channel_pool.session_management_channel
-        )
-        grpc_device_channel = grpc_channel_pool.get_grpc_device_channel(
-            nidmm.GRPC_SERVICE_INTERFACE_NAME
+        discovery_client = DiscoveryClient(grpc_channel_pool=grpc_channel_pool)
+        session_management_client = SessionManagementClient(
+            discovery_client=discovery_client, grpc_channel_pool=grpc_channel_pool
         )
         with session_management_client.reserve_all_registered_sessions(
-            instrument_type_id=nims.session_management.INSTRUMENT_TYPE_NI_DMM,
+            instrument_type_id=INSTRUMENT_TYPE_NI_DMM,
         ) as reservation:
+            if not reservation.session_info:
+                return
+
             session_management_client.unregister_sessions(reservation.session_info)
-            for session_info in reservation.session_info:
-                session = create_session(
-                    session_info,
-                    grpc_device_channel,
-                    initialization_behavior=nidmm.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
-                )
-                session.close()
+            with reservation.initialize_nidmm_sessions(
+                initialization_behavior=SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE
+            ):
+                pass

ni_measurementlink_service/_drivers/_nidaqmx.py

@@ -18,6 +18,8 @@
     SessionInitializationBehavior.AUTO: nidaqmx.SessionInitializationBehavior.AUTO,
     SessionInitializationBehavior.INITIALIZE_SERVER_SESSION: nidaqmx.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
     SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION: nidaqmx.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
+    SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH: nidaqmx.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
+    SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE: nidaqmx.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
 }
 
 

ni_measurementlink_service/_drivers/_nidcpower.py

@@ -19,6 +19,8 @@
     SessionInitializationBehavior.AUTO: nidcpower.SessionInitializationBehavior.AUTO,
     SessionInitializationBehavior.INITIALIZE_SERVER_SESSION: nidcpower.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
     SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION: nidcpower.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
+    SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH: nidcpower.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
+    SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE: nidcpower.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
 }
 
 

ni_measurementlink_service/_drivers/_nidmm.py

@@ -19,6 +19,8 @@
     SessionInitializationBehavior.AUTO: nidmm.SessionInitializationBehavior.AUTO,
     SessionInitializationBehavior.INITIALIZE_SERVER_SESSION: nidmm.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
     SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION: nidmm.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
+    SessionInitializationBehavior.INITIALIZE_SESSION_THEN_DETACH: nidmm.SessionInitializationBehavior.INITIALIZE_SERVER_SESSION,
+    SessionInitializationBehavior.ATTACH_TO_SESSION_THEN_CLOSE: nidmm.SessionInitializationBehavior.ATTACH_TO_SERVER_SESSION,
 }
 
 

ni_measurementlink_service/session_management/_reservation.py

@@ -611,7 +611,12 @@ def create_nidaqmx_task(
         session_constructor = SessionConstructor(
             self._discovery_client, self._grpc_channel_pool, initialization_behavior
         )
-        return self._initialize_session_core(session_constructor, INSTRUMENT_TYPE_NI_DAQMX)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_session_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DAQMX, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
     def create_nidaqmx_tasks(
@@ -643,7 +648,12 @@ def create_nidaqmx_tasks(
         session_constructor = SessionConstructor(
             self._discovery_client, self._grpc_channel_pool, initialization_behavior
         )
-        return self._initialize_sessions_core(session_constructor, INSTRUMENT_TYPE_NI_DAQMX)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_sessions_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DAQMX, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
     def get_nidaqmx_connection(
@@ -738,10 +748,15 @@ def initialize_nidcpower_session(
         session_constructor = SessionConstructor(
             self._discovery_client, self._grpc_channel_pool, reset, options, initialization_behavior
         )
-        return self._initialize_session_core(session_constructor, INSTRUMENT_TYPE_NI_DCPOWER)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_session_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DCPOWER, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
-    def create_nidcpower_sessions(
+    def initialize_nidcpower_sessions(
         self,
         reset: bool = False,
         options: Optional[Dict[str, Any]] = None,
@@ -778,7 +793,12 @@ def create_nidcpower_sessions(
         session_constructor = SessionConstructor(
             self._discovery_client, self._grpc_channel_pool, reset, options, initialization_behavior
         )
-        return self._initialize_sessions_core(session_constructor, INSTRUMENT_TYPE_NI_DCPOWER)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_sessions_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DCPOWER, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
     def get_nidcpower_connection(
@@ -1032,7 +1052,12 @@ def initialize_nidmm_session(
             options,
             initialization_behavior,
         )
-        return self._initialize_session_core(session_constructor, INSTRUMENT_TYPE_NI_DMM)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_session_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DMM, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
     def initialize_nidmm_sessions(
@@ -1076,7 +1101,12 @@ def initialize_nidmm_sessions(
             options,
             initialization_behavior,
         )
-        return self._initialize_sessions_core(session_constructor, INSTRUMENT_TYPE_NI_DMM)
+        closing_function = functools.partial(
+            closing_session_with_ts_code_module_support, initialization_behavior
+        )
+        return self._initialize_sessions_core(
+            session_constructor, INSTRUMENT_TYPE_NI_DMM, closing_function
+        )
 
     @requires_feature(SESSION_MANAGEMENT_2024Q1)
     def get_nidmm_connection(