Replaced magic strings with enums.

Signed-off-by: Sathish Sekar <[email protected]>

src/pci_lmt/collector.py

@@ -10,7 +10,7 @@
 import typing as ty
 
 from pci_lmt import __version__ as PCI_LMT_VERSION
-from pci_lmt.config import PlatformConfig
+from pci_lmt.config import MarginType, PlatformConfig
 from pci_lmt.device import PciDevice
 from pci_lmt.host import HostInfo
 from pci_lmt.pcie_lane_margining import PcieDeviceLaneMargining
@@ -27,8 +27,6 @@ def __init__(self, test_info: LmtTestInfo, bdf_list: ty.List[str]):
         # will handle this in a separate PR.
         self.receiver_number = test_info.receiver_number
         self.error_count_limit = test_info.error_count_limit
-        self.left_right_none = test_info.left_right_none
-        self.up_down = test_info.up_down
         self.step = test_info.step
         self.margin_type = test_info.margin_type
         self.force_margin = test_info.force_margin
@@ -97,39 +95,24 @@ def setup_lane_margin_on_device_list(self):
                     error_count_limit=self.error_count_limit,
                 )
 
-                if self.margin_type == "TIMING":
+                if self.margin_type in (MarginType.TIMING_LEFT, MarginType.TIMING_RIGHT, MarginType.TIMING_NONE):
                     dev.step_margin_timing_offset_right_left_of_default(
                         lane=lane,
                         receiver_number=self.receiver_number,
-                        left_right_none=self.left_right_none,
+                        margin_type=self.margin_type,
                         steps=self.step,
                     )
                 else:
                     dev.step_margin_voltage_offset_up_down_of_default(
                         lane=lane,
                         receiver_number=self.receiver_number,
-                        up_down=self.up_down,
+                        margin_type=self.margin_type,
                         steps=self.step,
                     )
 
-    # FIXME: `margin_type` should be an enum; dont use strings for magic constants
     def collect_lane_margin_on_device_list(self) -> ty.List[LmtLaneResult]:
         """Returns the Lane Margining Test result from all lanes as a list."""
 
-        def get_timing_margin_type_as_str() -> str:
-            if self.left_right_none == 0:
-                return "timing_right"
-            if self.left_right_none == 1:
-                return "timing_left"
-            return "timing_none"
-
-        def get_voltage_margin_type_as_str() -> str:
-            if self.up_down == 0:
-                return "voltage_up"
-            if self.up_down == 1:
-                return "voltage_down"
-            return "voltage_none"
-
         results = []
         for dev in self.devices:
             # Collect results from all devices and lanes
@@ -140,13 +123,12 @@ def get_voltage_margin_type_as_str() -> str:
                     lane=lane,
                     receiver_number=ty.cast(int, self.receiver_number),
                     step=self.step,
+                    margin_type=self.margin_type.value,
                 )
 
-                if self.margin_type == "TIMING":
-                    lane_result.margin_type = get_timing_margin_type_as_str()
+                if self.margin_type in (MarginType.TIMING_LEFT, MarginType.TIMING_RIGHT, MarginType.TIMING_NONE):
                     margin_status = dev.decode_step_margin_timing_offset_right_left_of_default(lane=lane)
                 else:
-                    lane_result.margin_type = get_voltage_margin_type_as_str()
                     margin_status = dev.decode_step_margin_voltage_offset_up_down_of_default(lane=lane)
 
                 sampler = dev.fetch_sample_count(lane=lane, receiver_number=self.receiver_number)
@@ -223,7 +205,6 @@ def run_lmt(args: argparse.Namespace, config: PlatformConfig, host: HostInfo, re
             test_info.margin_type = group.margin_type
             test_info.receiver_number = group.receiver_number
             test_info.annotation = args.annotation if args.annotation else group.name
-            test_info.left_right_none, test_info.up_down = group.margin_directions_tuple
 
             # Loop through each step running LMT on all BDFs.
             for step in group.margin_steps:

src/pci_lmt/config.py

@@ -8,44 +8,26 @@
 import textwrap
 import typing as ty
 
+from enum import Enum
+
+
+class MarginType(str, Enum):
+    VOLTAGE_NONE = "voltage_none"
+    VOLTAGE_UP = "voltage_up"
+    VOLTAGE_DOWN = "voltage_down"
+    TIMING_NONE = "timing_none"
+    TIMING_RIGHT = "timing_right"
+    TIMING_LEFT = "timing_left"
+
 
 @dc.dataclass
 class ConfigLmtGroup:
     name: str
     receiver_number: int
     bdf_list: ty.List[str]
-    margin_type: ty.Union[ty.Literal["VOLTAGE"], ty.Literal["TIMING"]]
-    margin_direction: ty.Union[
-        ty.Literal["right"],
-        ty.Literal["left"],
-        ty.Literal["up"],
-        ty.Literal["down"],
-    ]
+    margin_type: MarginType
     margin_steps: ty.List[int]
 
-    # NOTE: this may need to be refactored to return some enum types; a tuple of ints
-    # is not really descriptive
-    @property
-    def margin_directions_tuple(self) -> ty.Tuple[int, int]:
-        """Returns the margin direction as a tuple."""
-        left_right_none = -1
-        up_down = -1
-        margin_info = (self.margin_type, self.margin_direction)
-        if margin_info == ("TIMING", "right"):
-            left_right_none = 0
-        elif margin_info == ("TIMING", "left"):
-            left_right_none = 1
-        elif margin_info == ("VOLTAGE", "up"):
-            up_down = 0
-        elif margin_info == ("VOLTAGE", "down"):
-            up_down = 1
-        else:
-            raise ValueError(
-                f"Invalid values for margin_type {self.margin_type} and/or margin_direction {self.margin_direction}."
-            )
-
-        return (left_right_none, up_down)
-
     def __str__(self) -> str:
         bdf = textwrap.indent("\n".join(self.bdf_list), " " * 16).lstrip()
         return textwrap.dedent(
@@ -55,7 +37,6 @@ def __str__(self) -> str:
             bdf:
                 {bdf}
             type: {self.margin_type}
-            direction: {self.margin_direction}
             steps: {self.margin_steps}
         """
         )
@@ -66,8 +47,7 @@ def from_json(data: ty.Dict[str, ty.Any]) -> "ConfigLmtGroup":
             name=data["name"],
             receiver_number=data["receiver_number"],
             bdf_list=data["bdf_list"],
-            margin_type=data["margin_type"],
-            margin_direction=data["margin_direction"],
+            margin_type=MarginType(f"{data['margin_type'].lower()}_{data['margin_direction'].lower()}"),
             margin_steps=data["margin_steps"],
         )
 

src/pci_lmt/pcie_lane_margining.py

@@ -11,6 +11,7 @@
 
 from pci_lmt.constants import MARGIN_RESPONSE
 from pci_lmt.device import PciDevice
+from pci_lmt.config import MarginType
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -603,7 +604,7 @@ def clear_error_log(self, lane: int, receiver_number: int):
 
     @handle_lane_status
     def step_margin_timing_offset_right_left_of_default(
-        self, lane: int, receiver_number: int, left_right_none: int = 0, steps: int = 6
+        self, lane: int, receiver_number: int, margin_type: MarginType, steps: int = 6
     ):
         # LeftRightNone
         # 0 indicates to move the Receiver to the right of the normal setting to be used when ind_left_right_timing = 1.
@@ -632,17 +633,19 @@ def step_margin_timing_offset_right_left_of_default(
             return {"error": f"ERROR: StepMarginTimingOffsetRightLeftOfDefault - BAD receiver_number {receiver_number}"}
 
         if not self.device_info.ind_left_right_timing:
-            if left_right_none == -1:
+            if margin_type == MarginType.TIMING_NONE:
                 margin_payload = steps
             else:
                 return {
                     "error": "ERROR: StepMarginTimingOffsetRightLeftOfDefault - "
                     "Rcvr doesn't support independent left/right margining"
                 }
-        elif left_right_none in (0, 1):
-            margin_payload = left_right_none << 6 | steps
+        elif margin_type == MarginType.TIMING_RIGHT:
+            margin_payload = steps
+        elif margin_type == MarginType.TIMING_LEFT:
+            margin_payload = 1 << 6 | steps
         else:
-            return {"error": f"ERROR: StepMarginTimingOffsetRightLeftOfDefault - BAD left_right_none {left_right_none}"}
+            return {"error": f"ERROR: StepMarginTimingOffsetRightLeftOfDefault - BAD margin_type {margin_type}"}
 
         self.no_command(lane=lane)
         self.write_margining_lane_control_register(
@@ -707,7 +710,7 @@ def decode_step_margin_timing_offset_right_left_of_default(self, lane: int):
 
     @handle_lane_status
     def step_margin_voltage_offset_up_down_of_default(
-        self, lane: int, receiver_number: int, up_down: int = 0, steps: int = 32
+        self, lane: int, receiver_number: int, margin_type: MarginType, steps: int = 32
     ):
         # UpDown
         # 0 indicates to move the Receiver to the Up from Normal.
@@ -734,17 +737,19 @@ def step_margin_voltage_offset_up_down_of_default(
             return {"error": f"ERROR: StepMarginVoltageOffsetUpDownOfDefault - BAD receiver_number {receiver_number}"}
 
         if not self.device_info.ind_up_down_voltage:
-            if up_down == -1:
+            if margin_type == MarginType.VOLTAGE_NONE:
                 margin_payload = steps
             else:
                 return {
                     "error": "ERROR: StepMarginVoltageOffsetUpDownOfDefault - "
                     "Rcvr doesn't support independent up/down margining"
                 }
-        elif up_down in (0, 1):
-            margin_payload = up_down << 6 | steps
+        elif margin_type == MarginType.VOLTAGE_UP:
+            margin_payload = steps
+        elif margin_type == MarginType.VOLTAGE_DOWN:
+            margin_payload = 1 << 6 | steps
         else:
-            return {"error": f"ERROR: StepMarginVoltageOffsetUpDownOfDefault - BAD UpDown {up_down}"}
+            return {"error": f"ERROR: StepMarginVoltageOffsetUpDownOfDefault - BAD margin_type {margin_type}"}
 
         self.no_command(lane=lane)
         self.write_margining_lane_control_register(

src/pci_lmt/results.py

@@ -13,6 +13,7 @@
 from pci_lmt import __version__ as PCI_LMT_VERSION
 from pci_lmt.host import HostInfo
 from pci_lmt.pcie_lane_margining import LmtDeviceInfo
+from pci_lmt.config import MarginType
 
 
 @dc.dataclass
@@ -25,9 +26,7 @@ class LmtTestInfo:  # pylint: disable=too-many-instance-attributes,too-few-publi
     hostname: str = ""
     model_name: str = ""
     receiver_number: int = -1
-    margin_type: str = ""
-    left_right_none: int = -1
-    up_down: int = -1
+    margin_type: MarginType = MarginType.VOLTAGE_NONE
     step: int = -1
     force_margin: bool = False
     dwell_time_secs: int = -1