Merge pull request #55 from lsst-sqre/tickets/DM-30752B

Add stop() method to business class and provide defaults for implemen…

src/mobu/business.py

@@ -24,5 +24,8 @@ async def run(self) -> None:
             logger.info("Idling...")
             await asyncio.sleep(5)
 
+    async def stop(self) -> None:
+        pass
+
     def dump(self) -> dict:
         return {"name": "Idle"}

src/mobu/jupyterloginloop.py

@@ -45,6 +45,9 @@ async def run(self) -> None:
             self.failure_count += 1
             raise
 
+    async def stop(self) -> None:
+        await self._client.delete_lab()
+
     def dump(self) -> dict:
         return {
             "name": "JupyterLoginLoop",

src/mobu/jupyterpythonloop.py

@@ -8,19 +8,22 @@
 ]
 
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from mobu.business import Business
 from mobu.jupyterclient import JupyterClient
 
 
 @dataclass
 class JupyterPythonLoop(Business):
+    _client: JupyterClient = field(init=False)
+
     async def run(self) -> None:
         logger = self.monkey.log
         logger.info("Starting up...")
 
         client = JupyterClient(self.monkey.user, logger, self.options)
+        self._client = client
         await client.hub_login()
         await client.ensure_lab()
 
@@ -31,5 +34,8 @@ async def run(self) -> None:
             logger.info(reply)
             await asyncio.sleep(60)
 
+    async def stop(self) -> None:
+        await self._client.delete_lab()
+
     def dump(self) -> dict:
         return {"name": "JupyterPythonLoop"}

src/mobu/monkey.py

@@ -104,6 +104,7 @@ async def _runner(self) -> None:
             except asyncio.CancelledError:
                 self.log.info("Shutting down")
                 run = False
+                await self.business.stop()
             except Exception as e:
                 self.state = "ERROR"
                 self.log.exception(
@@ -116,6 +117,7 @@ async def _runner(self) -> None:
                 await asyncio.sleep(60)
 
     async def stop(self) -> None:
+        await self.business.stop()
         try:
             await self._job.close(timeout=0)
         except asyncio.TimeoutError:

src/mobu/notebookrunner.py

@@ -104,6 +104,9 @@ async def run(self) -> None:
                 f"```{self.code}``` generated: ```{e}```"
             )
 
+    async def stop(self) -> None:
+        await self._client.delete_lab()
+
     def dump(self) -> dict:
         return {
             "name": "NotebookRunner",

src/mobu/querymonkey.py

@@ -2,6 +2,7 @@
 import os
 import random
 import time
+from dataclasses import field
 
 import jinja2
 import pyvo
@@ -32,8 +33,9 @@ def generate_parameters() -> dict:
 class QueryMonkey(Business):
     success_count: int = 0
     failure_count: int = 0
+    _client: pyvo.dal.TAPService = field(init=False)
 
-    def _client(self) -> pyvo.dal.TAPService:
+    def _make_client(self) -> pyvo.dal.TAPService:
         tap_url = Configuration.environment_url + "/api/tap"
 
         s = requests.Session()
@@ -64,15 +66,15 @@ async def run(self) -> None:
                 "Query templates to choose from: %s", env.list_templates()
             )
 
-            service = self._client()
+            self._client = self._make_client()
 
             while True:
                 template_name = random.choice(env.list_templates())
                 template = env.get_template(template_name)
                 query = template.render(generate_parameters())
                 logger.info("Running: %s", query)
                 start = time.time()
-                await loop.run_in_executor(None, service.search, query)
+                await loop.run_in_executor(None, self._client.search, query)
                 end = time.time()
                 logger.info("Finished, took: %i seconds", end - start)
                 self.success_count += 1
@@ -81,6 +83,11 @@ async def run(self) -> None:
             self.failure_count += 1
             raise
 
+    async def stop(self) -> None:
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self._client.abort)
+        await loop.run_in_executor(None, self._client.delete)
+
     def dump(self) -> dict:
         return {
             "name": "QueryMonkey",