Fix bugs and apply improvements across client library (#197)

- Fix PyatsException to inherit from VirlException instead of Exception
- Fix ControllerNotFound to use __init__ instead of class attribute for message
- Fix ControllerNotFound instantiation (raise ControllerNotFound())
- Fix sync_states staleness handling for nodes, interfaces, and links
- Fix _import_link parameter order (iface_a_id before iface_b_id)
- Fix incorrect log message in Interface._set_interface_property
- Refactor Annotation validation using _VALID_KEYS frozensets
- Refactor Version comparisons using _as_tuple helper
- Refactor get_lab_list to pass params directly to session.get
- Optimize Node.next_available_interface with direct slicing

Author: tmikuska

virl2_client/exceptions.py

@@ -124,7 +124,7 @@ class MethodNotActive(VirlException):
     pass
 
 
-class PyatsException(Exception):
+class PyatsException(VirlException):
     """Base exception for pyATS integration errors."""
 
     pass
@@ -145,7 +145,8 @@ class PyatsDeviceNotFound(PyatsException):
 class ControllerNotFound(VirlException):
     """Raised when no CML controller node is found in the topology."""
 
-    message = "Controller not found"
+    def __init__(self) -> None:
+        super().__init__("Controller not found")
 
 
 class APIError(VirlException, httpx.HTTPStatusError):

virl2_client/models/annotation.py

@@ -21,7 +21,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 from ..exceptions import InvalidProperty
 from ..utils import check_stale, get_url_from_template, locked
@@ -33,13 +33,13 @@
     from .lab import Lab
 
     AnnotationTypeString = Literal["text", "line", "ellipse", "rectangle"]
-    AnnotationType = (
-        "Annotation"
-        | "AnnotationRectangle"
-        | "AnnotationEllipse"
-        | "AnnotationLine"
-        | "AnnotationText"
-    )
+    AnnotationType = Union[
+        "Annotation",
+        "AnnotationRectangle",
+        "AnnotationEllipse",
+        "AnnotationLine",
+        "AnnotationText",
+    ]
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -123,6 +123,20 @@ class Annotation:
         "annotation": "labs/{lab_id}/annotations/{annotation_id}",
     }
 
+    _VALID_KEYS: frozenset[str] = frozenset(
+        {
+            "id",
+            "type",
+            "border_color",
+            "border_style",
+            "color",
+            "thickness",
+            "x1",
+            "y1",
+            "z_index",
+        }
+    )
+
     def __init__(
         self,
         lab: Lab,
@@ -387,12 +401,14 @@ def is_valid_property(
         :param _property: The property name to validate.
         :returns: True if the property is valid for the given type, False otherwise.
         """
-        try:
-            assert annotation_type in _ANNOTATION_TYPES
-            assert _property in ANNOTATION_PROPERTY_MAP
-        except AssertionError:
+        if (
+            annotation_type not in _ANNOTATION_TYPES
+            or _property not in ANNOTATION_PROPERTY_MAP
+        ):
             return False
-        return ANNOTATION_MAP[annotation_type] & ANNOTATION_PROPERTY_MAP[_property] > 0
+        return (
+            ANNOTATION_MAP[annotation_type] & ANNOTATION_PROPERTY_MAP[_property]
+        ) > 0
 
     @locked
     def as_dict(self) -> dict[str, Any]:
@@ -449,9 +465,8 @@ def _update(self, annotation_data: dict[str, Any], push_to_server: bool) -> None
             raise ValueError("Can't change annotation type.")
 
         # make sure all properties we want to update are valid
-        existing_keys = dir(self)
         for key in annotation_data:
-            if key not in existing_keys:
+            if key not in self._VALID_KEYS:
                 raise InvalidProperty(f"Invalid annotation property: {key}")
 
         if push_to_server:
@@ -493,6 +508,10 @@ class AnnotationRectangle(Annotation):
     Annotation class representing rectangle annotation.
     """
 
+    _VALID_KEYS = Annotation._VALID_KEYS | frozenset(
+        {"border_radius", "x2", "y2", "rotation"}
+    )
+
     def __init__(
         self,
         lab: Lab,
@@ -600,6 +619,8 @@ class AnnotationEllipse(Annotation):
     Annotation class representing ellipse annotation.
     """
 
+    _VALID_KEYS = Annotation._VALID_KEYS | frozenset({"x2", "y2", "rotation"})
+
     def __init__(
         self,
         lab: Lab,
@@ -687,6 +708,10 @@ class AnnotationLine(Annotation):
     Annotation class representing line annotation.
     """
 
+    _VALID_KEYS = Annotation._VALID_KEYS | frozenset(
+        {"x2", "y2", "line_start", "line_end"}
+    )
+
     def __init__(
         self,
         lab: Lab,
@@ -794,6 +819,20 @@ class AnnotationText(Annotation):
     Annotation class representing text annotation.
     """
 
+    _VALID_KEYS = Annotation._VALID_KEYS | frozenset(
+        {
+            "x2",
+            "y2",
+            "rotation",
+            "text_bold",
+            "text_content",
+            "text_font",
+            "text_italic",
+            "text_size",
+            "text_unit",
+        }
+    )
+
     def __init__(
         self,
         lab: Lab,

virl2_client/models/interface.py

@@ -437,7 +437,7 @@ def _set_interface_property(self, key: str, val: Any) -> None:
         :param key: The key of the property to set.
         :param val: The value to set.
         """
-        _LOGGER.debug(f"Setting node property {self} {key}: {val}")
+        _LOGGER.debug(f"Setting interface property {self} {key}: {val}")
         self._set_interface_properties({key: val})
 
     @check_stale

virl2_client/models/lab.py

@@ -1357,8 +1357,18 @@ def sync_states(self) -> None:
         """
         url = self._url_for("lab_element_state")
         states: dict[str, dict[str, str]] = self._session.get(url).json()
+
+        nodes = self._nodes.copy()
         for node_id, node_state in states["nodes"].items():
-            self._nodes[node_id]._state = node_state
+            try:
+                node = nodes.pop(node_id)
+            except KeyError:
+                pass
+            else:
+                node._state = node_state
+        for stale_node in nodes.values():
+            stale_node._stale = True
+
         ifaces = self._interfaces.copy()
         for interface_id, interface_state in states["interfaces"].items():
             try:
@@ -1367,10 +1377,19 @@ def sync_states(self) -> None:
                 pass
             else:
                 iface._state = interface_state
-        for stale_iface in ifaces:
-            ifaces[stale_iface]._stale = True
+        for stale_iface in ifaces.values():
+            stale_iface._stale = True
+
+        links = self._links.copy()
         for link_id, link_state in states["links"].items():
-            self._links[link_id]._state = link_state
+            try:
+                link = links.pop(link_id)
+            except KeyError:
+                pass
+            else:
+                link._state = link_state
+        for stale_link in links.values():
+            stale_link._stale = True
 
         self._last_sync_state_time = time.time()
 
@@ -1714,7 +1733,7 @@ def _handle_import_links(self, topology: dict) -> None:
             iface_b_id = link["interface_b"]
             label = link.get("label")
 
-            self._import_link(link_id, iface_b_id, iface_a_id, label)
+            self._import_link(link_id, iface_a_id, iface_b_id, label)
 
     @locked
     def _handle_import_annotations(self, topology: dict) -> None:
@@ -1749,16 +1768,16 @@ def _handle_import_annotations(self, topology: dict) -> None:
     def _import_link(
         self,
         link_id: str,
-        iface_b_id: str,
         iface_a_id: str,
+        iface_b_id: str,
         label: str | None = None,
     ) -> Link:
         """
         Import a link with the given parameters.
 
         :param link_id: The ID of the link.
-        :param iface_b_id: The ID of the second interface.
         :param iface_a_id: The ID of the first interface.
+        :param iface_b_id: The ID of the second interface.
         :param label: The label of the link.
         :returns: The imported Link object.
         """
@@ -2069,7 +2088,7 @@ def _add_links(self, topology: dict, new_links: Iterable[str]) -> None:
             iface_a_id = link_data["interface_a"]
             iface_b_id = link_data["interface_b"]
             label = link_data.get("label")
-            link = self._import_link(link_id, iface_b_id, iface_a_id, label)
+            link = self._import_link(link_id, iface_a_id, iface_b_id, label)
             _LOGGER.info(f"Added link {link}")
 
     @locked

virl2_client/models/node.py

@@ -300,7 +300,7 @@ def next_available_interface(self, index: int = 0) -> Interface | None:
         :returns: An available physical interface or None if all existing
             ones are connected.
         """
-        for _, iface in enumerate(self.interfaces(), index):
+        for iface in self.interfaces()[index:]:
             if not iface.connected and iface.physical:
                 return iface
         return None

virl2_client/models/system.py

@@ -101,7 +101,7 @@ def controller(self) -> ComputeHost:
         for compute_host in self._compute_hosts.values():
             if compute_host.is_connector:
                 return compute_host
-        raise ControllerNotFound
+        raise ControllerNotFound()
 
     @property
     def system_notices(self) -> dict[str, SystemNotice]:

virl2_client/virl2_client.py

@@ -84,43 +84,33 @@ def parse_version_str(version_str: str) -> tuple[str, str, str, str]:
     def __repr__(self) -> str:
         return self.version_str
 
-    def __eq__(self, other):
-        return (
-            isinstance(other, self.__class__)
-            and self.major == other.major
-            and self.minor == other.minor
-            and self.patch == other.patch
-        )
-
-    def __gt__(self, other):
-        if isinstance(other, self.__class__):
-            if self.major > other.major:
-                return True
-            elif self.major == other.major:
-                if self.minor > other.minor:
-                    return True
-                elif self.minor == other.minor:
-                    if self.patch > other.patch:
-                        return True
-        return False
-
-    def __ge__(self, other):
-        return self == other or self > other
-
-    def __lt__(self, other):
-        if isinstance(other, self.__class__):
-            if self.major < other.major:
-                return True
-            elif self.major == other.major:
-                if self.minor < other.minor:
-                    return True
-                elif self.minor == other.minor:
-                    if self.patch < other.patch:
-                        return True
-        return False
-
-    def __le__(self, other):
-        return self == other or self < other
+    def _as_tuple(self) -> tuple[int, int, int]:
+        return (self.major, self.minor, self.patch)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return False
+        return self._as_tuple() == other._as_tuple()
+
+    def __gt__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return False
+        return self._as_tuple() > other._as_tuple()
+
+    def __ge__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return False
+        return self._as_tuple() >= other._as_tuple()
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return False
+        return self._as_tuple() < other._as_tuple()
+
+    def __le__(self, other: object) -> bool:
+        if not isinstance(other, Version):
+            return False
+        return self._as_tuple() <= other._as_tuple()
 
     def major_differs(self, other: Version) -> bool:
         return self.major != other.major
@@ -995,10 +985,8 @@ def get_lab_list(self, show_all: bool = False) -> list[str]:
             owned by the authenticated user (False).
         :returns: A list of lab IDs.
         """
-        url: dict[str, str | dict] = {"url": self._url_for("labs")}
-        if show_all:
-            url["params"] = {"show_all": True}
-        return self._session.get(**url).json()
+        params = {"show_all": True} if show_all else None
+        return self._session.get(self._url_for("labs"), params=params).json()
 
 
 def _prepare_url(url: str, allow_http: bool) -> tuple[str, str]: