Define shutdown method for NeonMQHandler with unit test coverage

NeonGeckoCom · Jan 14, 2025 · f6045fc · f6045fc
1 parent c5bd3b7
commit f6045fc
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/neon_mq_connector/utils/client_utils.py b/neon_mq_connector/utils/client_utils.py
@@ -57,6 +57,10 @@ def __init__(self, config: dict, service_name: str, vhost: str):
         self.connection = pika.BlockingConnection(
             parameters=self.get_connection_params(vhost))
 
+    def shutdown(self):
+        MQConnector.stop(self)
+        self.connection.close()
+
 
 def send_mq_request(vhost: str, request_data: dict, target_queue: str,
                     response_queue: str = None, timeout: int = 30,

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -43,7 +43,7 @@
 from neon_mq_connector.utils import RepeatingTimer
 from neon_mq_connector.utils.connection_utils import get_timeout, retry, \
     wait_for_mq_startup
-from neon_mq_connector.utils.client_utils import MQConnector
+from neon_mq_connector.utils.client_utils import MQConnector, NeonMQHandler
 from neon_mq_connector.utils.network_utils import dict_to_b64, b64_to_dict
 
 from .fixtures import rmq_instance
@@ -116,15 +116,15 @@ class TestClientUtils(unittest.TestCase):
 
     def setUp(self) -> None:
         if self.test_connector is None:
-            test_conf = {
+            self.test_conf = {
                 "server": "localhost",
                 "port": self.rmq_instance.port,
                 "users": {"mq_handler": {"user": "test_user",
                                          "password": "test_password"}}}
             import neon_mq_connector.utils.client_utils
-            neon_mq_connector.utils.client_utils._default_mq_config = test_conf
+            neon_mq_connector.utils.client_utils._default_mq_config = self.test_conf
             vhost = "/neon_testing"
-            self.test_connector = SimpleMQConnector(config=test_conf,
+            self.test_connector = SimpleMQConnector(config=self.test_conf,
                                                     service_name="mq_handler",
                                                     vhost=vhost)
             self.test_connector.register_consumer("neon_utils_test",
@@ -199,6 +199,14 @@ def test_send_mq_request_invalid_vhost(self):
         with self.assertRaises(ValueError):
             send_mq_request("invalid_endpoint", {}, "test", "test", timeout=5)
 
+    def test_connector_shutdown(self):
+        connector = NeonMQHandler(config=self.test_conf,
+                                  service_name="mq_handler",
+                                  vhost="/neon_testing")
+        self.assertTrue(connector.connection.is_open)
+        connector.shutdown()
+        self.assertTrue(connector.connection.is_closed)
+
 
 class TestMQConnectionUtils(unittest.TestCase):
     counter = 0