Update to gracefully close elastic

kytos/core/controller.py

@@ -154,6 +154,9 @@ def __init__(self, options=None, loop: AbstractEventLoop = None):
         self._alisten_tasks = set()
         self.qmonitors: list[QueueMonitorWindow] = []
 
+        #: APM client in memory to be closed when necessary
+        self.apm = None
+
         self._register_endpoints()
         #: Adding the napps 'enabled' directory into the PATH
         #: Now you can access the enabled napps with:
@@ -258,7 +261,7 @@ def start(self, restart=False):
                 db_conn_wait(db_backend=self.options.database)
                 self.start_auth()
             if self.options.apm:
-                init_apm(self.options.apm, app=self.api_server.app)
+                self.apm = init_apm(self.options.apm, app=self.api_server.app)
             if not restart:
                 self.create_pidfile()
             self.start_controller()
@@ -508,6 +511,9 @@ def stop_controller(self, graceful=True):
         # self.server.server_close()
 
         self.stop_queue_monitors()
+        if self.apm:
+            self.log.info("Stopping APM server...")
+            self.apm.close()
         self.log.info("Stopping API Server...")
         self.api_server.stop()
         self.log.info("Stopped API Server")

tests/unit/test_core/test_auth.py

@@ -1,4 +1,5 @@
 """Test kytos.core.auth module."""
+import asyncio
 import base64
 
 import pytest
@@ -97,9 +98,9 @@ async def test_02_list_users_request(self, auth, api_client):
         resp = await api_client.get(endpoint, headers=invalid_header)
         assert resp.status_code == 401
 
-    async def test_03_create_user_request(self, auth, api_client, event_loop):
+    async def test_03_create_user_request(self, auth, api_client):
         """Test auth create user endpoint."""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         endpoint = "kytos/core/auth/users"
         headers = await self.auth_headers(auth, api_client)
         resp = await api_client.post(endpoint, json=self.user_data,
@@ -108,10 +109,10 @@ async def test_03_create_user_request(self, auth, api_client, event_loop):
         assert "User successfully created" in resp.json()
 
     async def test_03_create_user_request_conflict(
-        self, auth, api_client, event_loop
+        self, auth, api_client
     ):
         """Test auth create user endpoint."""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         endpoint = "kytos/core/auth/users"
         auth.user_controller.create_user.side_effect = DuplicateKeyError("0")
         headers = await self.auth_headers(auth, api_client)
@@ -121,10 +122,10 @@ async def test_03_create_user_request_conflict(
         assert "already exists" in resp.json()["description"]
 
     async def test_03_create_user_request_bad(
-        self, auth, api_client, event_loop
+        self, auth, api_client
     ):
         """Test auth create user endpoint."""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         data = "wrong_json"
         endpoint = "kytos/core/auth/users"
         headers = await self.auth_headers(auth, api_client)
@@ -151,20 +152,19 @@ async def test_04_list_user_request_error(self, auth, api_client):
         assert resp.status_code == 404
         assert "not found" in resp.json()["description"]
 
-    async def test_05_update_user_request(self, auth, api_client, event_loop):
+    async def test_05_update_user_request(self, auth, api_client):
         """Test auth update user endpoint."""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         data = {"email": "[email protected]"}
         endpoint = f"kytos/core/auth/users/{self.user_data['username']}"
         headers = await self.auth_headers(auth, api_client)
         resp = await api_client.patch(endpoint, json=data,
                                       headers=headers)
         assert resp.status_code == 200
 
-    async def test_05_update_user_request_not_found(self, auth, api_client,
-                                                    event_loop):
+    async def test_05_update_user_request_not_found(self, auth, api_client):
         """Test auth update user endpoint."""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         data = {"email": "[email protected]"}
         auth.user_controller.update_user.return_value = {}
         endpoint = "kytos/core/auth/users/user5"
@@ -174,10 +174,9 @@ async def test_05_update_user_request_not_found(self, auth, api_client,
         assert resp.status_code == 404
         assert "not found" in resp.json()["description"]
 
-    async def test_05_update_user_request_bad(self, auth, api_client,
-                                              event_loop):
+    async def test_05_update_user_request_bad(self, auth, api_client):
         """Test auth update user endpoint"""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         exc = ValidationError.from_exception_data('', [])
         auth.user_controller.update_user.side_effect = exc
         endpoint = "kytos/core/auth/users/user5"
@@ -186,10 +185,9 @@ async def test_05_update_user_request_bad(self, auth, api_client,
                                       headers=headers)
         assert resp.status_code == 400
 
-    async def test_05_update_user_request_conflict(self, auth, api_client,
-                                                   event_loop):
+    async def test_05_update_user_request_conflict(self, auth, api_client):
         """Test auth update user endpoint"""
-        auth.controller.loop = event_loop
+        auth.controller.loop = asyncio.get_running_loop()
         auth.user_controller.update_user.side_effect = DuplicateKeyError("0")
         endpoint = "kytos/core/auth/users/user5"
         headers = await self.auth_headers(auth, api_client)

tests/unit/test_core/test_controller.py

@@ -182,6 +182,7 @@ def test_start(self, *args):
         mock_start_controller.assert_called()
         mock_db_conn_wait.assert_not_called()
         mock_init_apm.assert_not_called()
+        assert self.controller.apm is None
 
     @patch('kytos.core.controller.sys')
     @patch('kytos.core.controller.init_apm')
@@ -223,6 +224,7 @@ def test_start_with_mongodb_and_apm(self, *args):
         mock_start_controller.assert_called()
         mock_db_conn_wait.assert_called()
         mock_init_apm.assert_called()
+        assert self.controller.apm is not None
 
     @patch('kytos.core.controller.sys.exit')
     @patch('kytos.core.controller.Controller.create_pidfile')
@@ -749,9 +751,10 @@ async def test_stop_controller(self, controller):
         controller.api_server = api_server
         controller.napp_dir_listener = napp_dir_listener
         controller.stop_queue_monitors = MagicMock()
+        controller.apm = MagicMock()
 
         controller.stop_controller()
-
+        controller.apm.close.assert_called()
         controller.buffers.send_stop_signal.assert_called()
         api_server.stop.assert_called()
         napp_dir_listener.stop.assert_called()

tests/unit/test_core/test_rest_api_routes.py

@@ -61,10 +61,10 @@ def test_api_500_traceback_by_default() -> None:
     assert KytosConfig().options["daemon"].api_traceback_on_500
 
 
-async def test_get_json_or_400(controller, api_client, event_loop) -> None:
+async def test_get_json_or_400(controller, api_client) -> None:
     """Test get_json_or_400."""
 
-    controller.loop = event_loop
+    controller.loop = asyncio.get_running_loop()
 
     def handler(request: Request) -> JSONResponse:
         body = get_json_or_400(request, controller.loop)
@@ -98,9 +98,9 @@ async def test_error_msg():
     assert actual_msg == expected_msg
 
 
-async def test_get_body(controller, api_client, event_loop) -> None:
+async def test_get_body(controller, api_client) -> None:
     """Test get_body (low level-ish usage for validators)."""
-    controller.loop = event_loop
+    controller.loop = asyncio.get_running_loop()
 
     def handler(request: Request) -> JSONResponse:
         body_bytes = get_body(request, controller.loop)