fix remaining minor issues

Author: tmunzer-AIDE

src/mistapi/__api_session.py

@@ -856,7 +856,11 @@ def login_with_return(
             LOGGER.error("apisession:login_with_return:credentials are missing")
             return {"authenticated": False, "error": "credentials are missing"}
 
-        if resp.status_code == 200 and not resp.data.get("two_factor_required", False):
+        if (
+            resp.status_code == 200
+            and isinstance(resp.data, dict)
+            and not resp.data.get("two_factor_required", False)
+        ):
             LOGGER.info("apisession:login_with_return:access authorized")
             return {"authenticated": True, "error": ""}
         else:
@@ -884,7 +888,8 @@ def logout(self) -> None:
                 self._set_authenticated(False)
             else:
                 try:
-                    CONSOLE.error(resp.data["detail"])
+                    if isinstance(resp.data, dict) and "detail" in resp.data:
+                        CONSOLE.error(resp.data["detail"])
                 except (KeyError, TypeError, AttributeError):
                     if isinstance(resp.raw_data, bytes):
                         CONSOLE.error(resp.raw_data.decode("utf-8", errors="replace"))
@@ -1077,7 +1082,7 @@ def _getself(self) -> bool:
         uri = "/api/v1/self"
         LOGGER.info('apisession:_getself: sending GET request to "%s"', uri)
         resp = self.mist_get(uri)
-        if resp.status_code == 200 and resp.data:
+        if resp.status_code == 200 and resp.data and isinstance(resp.data, dict):
             # Deal with 2FA if needed
             if (
                 resp.data.get("two_factor_required") is True
@@ -1094,20 +1099,27 @@ def _getself(self) -> bool:
                 LOGGER.info(
                     "apisession:_getself:authentication Ok. Processing account privileges"
                 )
-                for key, val in resp.data.items():
-                    if key == "privileges":
-                        self.privileges = Privileges(resp.data["privileges"])
-                    if key == "tags":
-                        for tag in resp.data["tags"]:
-                            self.tags.append(tag)
-                    else:
-                        setattr(self, key, val)
-                if self._show_cli_notif:
-                    print()
-                    print(" Authenticated ".center(80, "-"))
-                    print(f"\r\nWelcome {self.first_name} {self.last_name}!\r\n")
-                LOGGER.info("apisession:_getself:account info processed successfully")
-                return True
+                if isinstance(resp.data, dict):
+                    for key, val in resp.data.items():
+                        if key == "privileges":
+                            self.privileges = Privileges(resp.data["privileges"])
+                        if key == "tags":
+                            for tag in resp.data["tags"]:
+                                self.tags.append(tag)
+                        else:
+                            setattr(self, key, val)
+                    if self._show_cli_notif:
+                        print()
+                        print(" Authenticated ".center(80, "-"))
+                        print(f"\r\nWelcome {self.first_name} {self.last_name}!\r\n")
+                    LOGGER.info(
+                        "apisession:_getself:account info processed successfully"
+                    )
+                    return True
+                else:
+                    raise ValueError(
+                        "Unexpected format for privileges in the response data"
+                    )
         elif resp.proxy_error:
             LOGGER.critical("apisession:_getself:proxy not valid...")
             CONSOLE.critical("Proxy not valid...\r\n")
@@ -1171,7 +1183,11 @@ def get_privilege_by_org_id(self, org_id: str):
             msp_id = None
             try:
                 resp = self.mist_get(uri)
-                if resp.data and resp.data.get("msp_id"):
+                if (
+                    resp.data
+                    and isinstance(resp.data, dict)
+                    and resp.data.get("msp_id")
+                ):
                     LOGGER.info(
                         "apisession:get_privilege_by_org_id:org %s belong to msp_id %s",
                         {org_id},
@@ -1205,7 +1221,7 @@ def get_privilege_by_org_id(self, org_id: str):
                         "unable of find msp %s privileges in user data",
                         msp_id,
                     )
-                else:
+                elif isinstance(resp.data, dict):
                     return {
                         "scope": "org",
                         "org_id": org_id,

tests/unit/test_websocket_client.py

@@ -9,7 +9,6 @@
 """
 
 import json
-import queue
 import ssl
 from unittest.mock import Mock, call, patch
 
@@ -25,19 +24,20 @@
 )
 from mistapi.websockets.orgs import (
     InsightsEvents,
-    MxEdgesStatsEvents as OrgMxEdgesStatsEvents,
-    MxEdgesUpgradesEvents,
+    MxEdgesEvents,
 )
+from mistapi.websockets.orgs import MxEdgesStatsEvents as OrgMxEdgesStatsEvents
 from mistapi.websockets.session import SessionWithUrl
 from mistapi.websockets.sites import (
     ClientsStatsEvents,
     DeviceCmdEvents,
+    DeviceEvents,
     DeviceStatsEvents,
-    DeviceUpgradesEvents,
-    MxEdgesStatsEvents as SiteMxEdgesStatsEvents,
     PcapEvents,
 )
-
+from mistapi.websockets.sites import (
+    MxEdgesStatsEvents as SiteMxEdgesStatsEvents,
+)
 
 # ---------------------------------------------------------------------------
 # Fixtures
@@ -694,7 +694,7 @@ def test_device_stats_events_channels(self, mock_session) -> None:
         assert ws._channels == ["/sites/s1/stats/devices"]
 
     def test_device_upgrades_events_channels(self, mock_session) -> None:
-        ws = DeviceUpgradesEvents(mock_session, site_ids=["s1"])
+        ws = DeviceEvents(mock_session, site_ids=["s1"])
         assert ws._channels == ["/sites/s1/devices"]
 
     def test_site_mxedges_stats_events_channels(self, mock_session) -> None:
@@ -729,7 +729,7 @@ def test_org_mxedges_stats_events_channels(self, mock_session) -> None:
         assert ws._channels == ["/orgs/o1/stats/mxedges"]
 
     def test_mxedges_upgrades_events_channels(self, mock_session) -> None:
-        ws = MxEdgesUpgradesEvents(mock_session, org_id="o1")
+        ws = MxEdgesEvents(mock_session, org_id="o1")
         assert ws._channels == ["/orgs/o1/mxedges"]
 
     def test_inherits_from_mist_websocket(self, mock_session) -> None: