diff --git a/src/nitsm/tsmcontext.py b/src/nitsm/tsmcontext.py index 2b04bf1..e2e98e9 100644 --- a/src/nitsm/tsmcontext.py +++ b/src/nitsm/tsmcontext.py @@ -1,4 +1,5 @@ """TSM Context Wrapper""" + import ctypes.wintypes import time import typing @@ -208,6 +209,24 @@ def site_numbers(self) -> "_Tuple[int, ...]": return self._context.SiteNumbers # Site and Global Data + def get_semiconductor_module_context_with_sites( + self, site_numbers: "_Sequence[int]" + ) -> "SemiconductorModuleContext": + """Returns a Semiconductor Module context object which holds information and resources + specific to the site_numbers mentioned. + + Args: + site_numbers: A sequence of site numbers for which the resources should be used. + + Returns: + SemiconductorModuleContext object with resources specific to the site_numbers. + """ + tsm_dispatch = self._context.GetSemiconductorModuleContextWithSites(site_numbers) + semiconductor_module_context_with_sites = SemiconductorModuleContext.__new__( + SemiconductorModuleContext + ) + semiconductor_module_context_with_sites._context = tsm_dispatch + return semiconductor_module_context_with_sites def set_site_data(self, data_id: str, data: "_Sequence[_Any]") -> None: """Associates a data item with each site. You can associate data with all sites or with the diff --git a/tests/multi_site.pinmap b/tests/multi_site.pinmap new file mode 100644 index 0000000..0b12940 --- /dev/null +++ b/tests/multi_site.pinmap @@ -0,0 +1,27 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/tests/test_multi_site.py b/tests/test_multi_site.py new file mode 100644 index 0000000..f9891ed --- /dev/null +++ b/tests/test_multi_site.py @@ -0,0 +1,34 @@ +import pytest +from nitsm.codemoduleapi import SemiconductorModuleContext + + +@pytest.mark.pin_map("multi_site.pinmap") +class TestMultiSite: + pin_map_dut_pins = ["DUTPin1"] + pin_map_system_pins = ["SystemPin1", "SystemPin2"] + sites = 4 # [0, 1, 2, 3] + + def test_get_semiconductor_module_with_sites(self, standalone_tsm_context): + # Get context for sites 1 and 3 + site_numbers = [1, 3] + filtered_tsm_context = standalone_tsm_context.get_semiconductor_module_context_with_sites( + site_numbers + ) + filtered_sites = list(filtered_tsm_context.site_numbers) + + # Validate the site numbers + assert len(site_numbers) == len(filtered_sites) + assert site_numbers == filtered_sites + for site in site_numbers: + assert site in filtered_sites + + # Validate that a fully initialized SemiconductorModuleContext wrapper is returned + assert isinstance(filtered_tsm_context, SemiconductorModuleContext) + # Exercise another wrapper method to catch initialization regressions + pin_names = list( + filtered_tsm_context.get_pin_names() + ) # Should contain DUTPins and SystemPins as a tuple + assert isinstance(pin_names, list) + # Each item is a tuple, change it to list and compare + assert list(pin_names[0]) == self.pin_map_dut_pins + assert list(pin_names[1]) == self.pin_map_system_pins