From c1084bc71e27014206719122c68ffc99db9a2289 Mon Sep 17 00:00:00 2001
From: adam <athornton@lsst.org>
Date: Mon, 14 Jun 2021 17:48:14 -0700
Subject: [PATCH] Add stop() method to business class and provide defaults for
 implementations

---
 src/mobu/business.py          |  3 +++
 src/mobu/jupyterloginloop.py  |  3 +++
 src/mobu/jupyterpythonloop.py |  8 +++++++-
 src/mobu/monkey.py            |  2 ++
 src/mobu/notebookrunner.py    |  3 +++
 src/mobu/querymonkey.py       | 13 ++++++++++---
 6 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/mobu/business.py b/src/mobu/business.py
index b7d7e105..c2117834 100644
--- a/src/mobu/business.py
+++ b/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"}
diff --git a/src/mobu/jupyterloginloop.py b/src/mobu/jupyterloginloop.py
index e9b72a25..a3e36aeb 100644
--- a/src/mobu/jupyterloginloop.py
+++ b/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",
diff --git a/src/mobu/jupyterpythonloop.py b/src/mobu/jupyterpythonloop.py
index 1b3e7050..e8a6fab5 100644
--- a/src/mobu/jupyterpythonloop.py
+++ b/src/mobu/jupyterpythonloop.py
@@ -8,7 +8,7 @@
 ]
 
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from mobu.business import Business
 from mobu.jupyterclient import JupyterClient
@@ -16,11 +16,14 @@
 
 @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"}
diff --git a/src/mobu/monkey.py b/src/mobu/monkey.py
index 3ee5045b..11200fcd 100644
--- a/src/mobu/monkey.py
+++ b/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:
diff --git a/src/mobu/notebookrunner.py b/src/mobu/notebookrunner.py
index 24353e7d..6a6c3d88 100644
--- a/src/mobu/notebookrunner.py
+++ b/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",
diff --git a/src/mobu/querymonkey.py b/src/mobu/querymonkey.py
index 8a804087..5589a449 100644
--- a/src/mobu/querymonkey.py
+++ b/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,7 +66,7 @@ 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())
@@ -72,7 +74,7 @@ async def run(self) -> None:
                 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",