From f13d7ad0976e6bab0d1d5d8b146c33b7da1fec4d Mon Sep 17 00:00:00 2001
From: Hanne Moa <hanne.moa@sikt.no>
Date: Wed, 3 Jul 2024 11:03:47 +0200
Subject: [PATCH] Improve _verify_session

* Disconnect safely, avoid raising an exception on disconnect
* Have _verify_session raise NotConnectedError instead of ValueError
* Update session fake in tests for the new golden path
* Add explicit tests for _verify_session
---
 src/zinolib/controllers/zino1.py        | 13 ++++++++++---
 tests/test_zinolib_controllers_zino1.py | 26 ++++++++++++++++++++++++-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/zinolib/controllers/zino1.py b/src/zinolib/controllers/zino1.py
index 235cd05..f2e3847 100644
--- a/src/zinolib/controllers/zino1.py
+++ b/src/zinolib/controllers/zino1.py
@@ -523,7 +523,11 @@ def _verify_session(self, quiet=False):
         if not getattr(self.session, 'request', None):
             if quiet:
                 return False
-            raise ValueError
+            raise NotConnectedError("The request socket have not been set up correctly. Reconnect necessary.")
+        if not self.session.request.connected:
+            if quiet:
+                return False
+            raise NotConnectedError("Authentication necessary")
         return True
 
     @classmethod
@@ -548,8 +552,11 @@ def authenticate(self, username=None, password=None):
             raise self.ManagerException(e)
 
     def disconnect(self):
-        self._verify_session()
-        self.session = self._session_adapter.close_session(self.session)
+        session_ok = self._verify_session(quiet=True)
+        if session_ok:
+            self.session = self._session_adapter.close_session(self.session)
+        else:
+            self._session_adapter.close_push_channel(self.session)
 
     def clear_flapping(self, event_or_id: EventOrId):
         """Clear flapping state of a PortStateEvent
diff --git a/tests/test_zinolib_controllers_zino1.py b/tests/test_zinolib_controllers_zino1.py
index 994c098..e56852a 100644
--- a/tests/test_zinolib_controllers_zino1.py
+++ b/tests/test_zinolib_controllers_zino1.py
@@ -2,7 +2,8 @@
 from datetime import datetime, timedelta, timezone
 
 from zinolib.event_types import AdmState, Event, HistoryEntry, LogEntry
-from zinolib.controllers.zino1 import EventAdapter, HistoryAdapter, LogAdapter, SessionAdapter, Zino1EventManager, UpdateHandler,RetryError
+from zinolib.controllers.zino1 import EventAdapter, HistoryAdapter, LogAdapter, SessionAdapter, Zino1EventManager, UpdateHandler
+from zinolib.controllers.zino1 import RetryError, NotConnectedError
 from zinolib.ritz import NotifierResponse
 
 raw_event_id = 139110
@@ -84,6 +85,7 @@ def _setup_config(cls, config):
     def _setup_request(session, config):
         class FakeSession:
             authenticated = True
+            connected = True
 
         session.request = FakeSession()  # needs to be truthy
         session.push = True  # needs to be truthy
@@ -106,6 +108,28 @@ def init_manager(self):
         zino1 = FakeZino1EventManager.configure(None)
         return zino1
 
+    def test_verify_session_raises_notconnectederror_on_incorrect_manager(self):
+        zino1 = FakeZino1EventManager()
+        with self.assertRaises(NotConnectedError) as e:
+            zino1._verify_session()
+            self.assertIn("The request socket have not been set up correctly", e)
+
+    def test_verify_session_raises_notconnectederror_if_not_connected(self):
+        zino1 = FakeZino1EventManager.configure(None)
+        orig = zino1.session.request.connected
+        zino1.session.request.connected = False
+        try:
+            with self.assertRaises(NotConnectedError) as e:
+                zino1._verify_session()
+                self.assertIn("Authentication necessary", e)
+        finally:
+            zino1.session.request.connected = orig
+
+    def test_verify_session_returns_False_if_quieted_on_incorrect_manager(self):
+        zino1 = FakeZino1EventManager()
+        result = zino1._verify_session(quiet=True)
+        self.assertEqual(result, False)
+
     def test_create_event_from_id_receiving_garbage_admstate_is_safely_handled(self):
         global raw_attrlist
         zino1 = self.init_manager()