From 7ec30451301f9ab4cb9a34ee79ae508079201082 Mon Sep 17 00:00:00 2001
From: Vlad Stan <stan.v.vlad@gmail.com>
Date: Fri, 17 Feb 2023 14:57:31 +0200
Subject: [PATCH] feat: add extension clean-up endpoint

---
 __init__.py             |  5 ++++-
 relay/client_manager.py |  4 ++++
 tests/test_clients.py   |  8 ++++----
 views_api.py            | 18 +++++++++++++++++-
 4 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/__init__.py b/__init__.py
index 4067bb4..b5d8407 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,4 +1,5 @@
 import asyncio
+from typing import List
 
 from fastapi import APIRouter
 from fastapi.staticfiles import StaticFiles
@@ -19,6 +20,7 @@ nostrrelay_static_files = [
     }
 ]
 
+scheduled_tasks: List[asyncio.Task] = []
 
 def nostrrelay_renderer():
     return template_renderer(["lnbits/extensions/nostrrelay/templates"])
@@ -31,4 +33,5 @@ from .views_api import *  # noqa
 
 def nostrrelay_start():
     loop = asyncio.get_event_loop()
-    loop.create_task(catch_everything_and_restart(wait_for_paid_invoices))
+    task = loop.create_task(catch_everything_and_restart(wait_for_paid_invoices))
+    scheduled_tasks.append(task)
diff --git a/relay/client_manager.py b/relay/client_manager.py
index c2db58d..488a719 100644
--- a/relay/client_manager.py
+++ b/relay/client_manager.py
@@ -52,6 +52,10 @@ class NostrClientManager:
             self._clients[relay_id] = []
         return self._clients[relay_id]
 
+    async def stop(self):
+        for relay_id in self._active_relays:
+            await self._stop_clients_for_relay(relay_id)
+
     async def _stop_clients_for_relay(self, relay_id: str):
         for client in self.clients(relay_id):
             if client.relay_id == relay_id:
diff --git a/tests/test_clients.py b/tests/test_clients.py
index e79c3bf..4035c45 100644
--- a/tests/test_clients.py
+++ b/tests/test_clients.py
@@ -6,11 +6,11 @@ import pytest
 from fastapi import WebSocket
 from loguru import logger
 
-from lnbits.extensions.nostrrelay.relay.client_connection import (
-    NostrClientConnection,  # type: ignore
+from lnbits.extensions.nostrrelay.relay.client_connection import (  # type: ignore
+    NostrClientConnection,
 )
-from lnbits.extensions.nostrrelay.relay.client_manager import (
-    NostrClientManager,  # type: ignore
+from lnbits.extensions.nostrrelay.relay.client_manager import (  # type: ignore
+    NostrClientManager,
 )
 from lnbits.extensions.nostrrelay.relay.relay import RelaySpec  # type: ignore
 
diff --git a/views_api.py b/views_api.py
index d30d523..2700378 100644
--- a/views_api.py
+++ b/views_api.py
@@ -15,7 +15,7 @@ from lnbits.decorators import (
 )
 from lnbits.helpers import urlsafe_short_hash
 
-from . import nostrrelay_ext
+from . import nostrrelay_ext, scheduled_tasks
 from .crud import (
     create_account,
     create_relay,
@@ -283,3 +283,19 @@ async def api_pay_to_join(data: BuyOrder):
             status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
             detail="Cannot create invoice for client to join",
         )
+
+
+@nostrrelay_ext.delete("/api/v1", status_code=HTTPStatus.OK)
+async def api_stop(wallet: WalletTypeInfo = Depends(check_admin)):
+    for t in scheduled_tasks:
+        try:
+            t.cancel()
+        except Exception as ex:
+            logger.warning(ex)
+
+    try:
+        await client_manager.stop()
+    except Exception as ex:
+        logger.warning(ex)
+
+    return {"success": True}