From 7b98b2332f6f675131b90e2f245e49764b535055 Mon Sep 17 00:00:00 2001
From: Luna <git@l4.pm>
Date: Sun, 5 Sep 2021 14:04:59 -0300
Subject: [PATCH] tests: add test resource awareness

---
 tests/common.py   | 53 +++++++++++++++++++++++++++++++++++++++++++++--
 tests/conftest.py |  8 +++++--
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/tests/common.py b/tests/common.py
index 4e573a7..ddb4ac2 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -18,24 +18,73 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 
 import secrets
+from typing import Optional
+from dataclasses import dataclass
+
+from litecord.common.users import create_user, delete_user
+from litecord.blueprints.auth import make_token
 
 
 def email() -> str:
     return f"{secrets.token_hex(5)}@{secrets.token_hex(5)}.com"
 
 
+@dataclass
+class WrappedUser:
+    test_cli: "TestClient"
+    id: int
+    name: str
+    email: str
+    password: str
+    token: str
+
+    async def refetch(self):
+        async with self.test_cli.app.app_context():
+            return await self.test_cli.app.storage.get_user(self.id)
+
+    async def delete(self):
+        async with self.test_cli.app.app_context():
+            return await delete_user(self.id)
+
+
 class TestClient:
-    """Test client that wraps pytest-sanic's TestClient and a test
-    user and adds authorization headers to test requests."""
+    """Test client wrapper class. Adds Authorization headers to all requests
+    and manages test resource setup and destruction."""
 
     def __init__(self, test_cli, test_user):
         self.cli = test_cli
         self.app = test_cli.app
         self.user = test_user
+        self.resources = []
 
     def __getitem__(self, key):
         return self.user[key]
 
+    def add_resource(self, resource):
+        self.resources.append(resource)
+        return resource
+
+    async def cleanup(self):
+        for resource in self.resources:
+            await resource.delete()
+
+    async def create_user(
+        self,
+        *,
+        username: str,
+        email: str,
+        password: Optional[str] = None,
+    ) -> WrappedUser:
+        password = password or secrets.token_hex(6)
+
+        async with self.app.app_context():
+            user_id, password_hash = await create_user(username, email, password)
+            user_token = make_token(user_id, password_hash)
+
+        return self.add_resource(
+            WrappedUser(self, user_id, username, email, password, user_token)
+        )
+
     def _inject_auth(self, kwargs: dict) -> list:
         """Inject the test user's API key into the test request before
         passing the request on to the underlying TestClient."""
diff --git a/tests/conftest.py b/tests/conftest.py
index f0444c9..57df793 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -107,7 +107,9 @@ async def test_user_fixture(app):
 async def test_cli_user(test_cli, test_user):
     """Yield a TestClient instance that contains a randomly generated
     user."""
-    yield TestClient(test_cli, test_user)
+    client = TestClient(test_cli, test_user)
+    yield client
+    await client.cleanup()
 
 
 @pytest.fixture
@@ -138,5 +140,7 @@ async def test_cli_staff(test_cli):
         user_id,
     )
 
-    yield TestClient(test_cli, test_user)
+    client = TestClient(test_cli, test_user)
+    yield client
+    await client.cleanup()
     await _user_fixture_teardown(test_cli.app, test_user)