diff --git a/examples/nidaqmx_analog_input/teststand_nidaqmx.py b/examples/nidaqmx_analog_input/teststand_nidaqmx.py index 6caedaa65..9ea32cf09 100644 --- a/examples/nidaqmx_analog_input/teststand_nidaqmx.py +++ b/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,27 +19,23 @@ 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) @@ -43,21 +43,18 @@ def create_nidaqmx_tasks(sequence_context: Any) -> None: 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 diff --git a/examples/nidcpower_source_dc_voltage/teststand_nidcpower.py b/examples/nidcpower_source_dc_voltage/teststand_nidcpower.py index 24bf71ed2..4273defee 100644 --- a/examples/nidcpower_source_dc_voltage/teststand_nidcpower.py +++ b/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 diff --git a/examples/nidmm_measurement/teststand_nidmm.py b/examples/nidmm_measurement/teststand_nidmm.py index 6e258bc97..f68b01784 100644 --- a/examples/nidmm_measurement/teststand_nidmm.py +++ b/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,27 +19,22 @@ 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) @@ -43,20 +42,18 @@ def create_nidmm_sessions(sequence_context: Any) -> None: 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 diff --git a/ni_measurementlink_service/_drivers/_nidaqmx.py b/ni_measurementlink_service/_drivers/_nidaqmx.py index cbac5c6a6..5136d73d3 100644 --- a/ni_measurementlink_service/_drivers/_nidaqmx.py +++ b/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, } diff --git a/ni_measurementlink_service/_drivers/_nidcpower.py b/ni_measurementlink_service/_drivers/_nidcpower.py index 7dd8c39f3..6c0cb2929 100644 --- a/ni_measurementlink_service/_drivers/_nidcpower.py +++ b/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, } diff --git a/ni_measurementlink_service/_drivers/_nidmm.py b/ni_measurementlink_service/_drivers/_nidmm.py index 7eb8835ee..d8a8666db 100644 --- a/ni_measurementlink_service/_drivers/_nidmm.py +++ b/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, } diff --git a/ni_measurementlink_service/session_management/_reservation.py b/ni_measurementlink_service/session_management/_reservation.py index b308a8df5..d9702c176 100644 --- a/ni_measurementlink_service/session_management/_reservation.py +++ b/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( diff --git a/tests/unit/_drivers/test_nidcpower.py b/tests/unit/_drivers/test_nidcpower.py index 0c33b5d54..a1c040d3c 100644 --- a/tests/unit/_drivers/test_nidcpower.py +++ b/tests/unit/_drivers/test_nidcpower.py @@ -62,7 +62,7 @@ def test___multiple_session_infos___initialize_nidcpower_sessions___sessions_cre sessions = create_mock_nidcpower_sessions(3) session_new.side_effect = sessions - with reservation.create_nidcpower_sessions() as session_info: + with reservation.initialize_nidcpower_sessions() as session_info: assert session_info[0].session == sessions[0] assert session_info[1].session == sessions[1] @@ -189,7 +189,7 @@ def test___session_created___get_nidcpower_connection___connection_returned( reservation = MultiSessionReservation(session_management_client, grpc_session_infos) sessions = create_mock_nidcpower_sessions(2) session_new.side_effect = sessions - session_infos = stack.enter_context(reservation.create_nidcpower_sessions()) + session_infos = stack.enter_context(reservation.initialize_nidcpower_sessions()) connection = reservation.get_nidcpower_connection(**kwargs) @@ -236,7 +236,7 @@ def test___session_created___get_nidcpower_connections___connections_returned( reservation = MultiSessionReservation(session_management_client, grpc_session_infos) sessions = create_mock_nidcpower_sessions(2) session_new.side_effect = sessions - session_infos = stack.enter_context(reservation.create_nidcpower_sessions()) + session_infos = stack.enter_context(reservation.initialize_nidcpower_sessions()) connections = reservation.get_nidcpower_connections(**kwargs)