add tests for vacuum

4 files changed, 106 insertions, 50 deletions

diff --git a/pypaste/server/__init__.py b/pypaste/server/__init__.py
index 5d2fd6d..da5e538 100644
--- a/pypaste/server/__init__.py
+++ b/pypaste/server/__init__.py
@@ -131,6 +131,14 @@ class Storage:
     async def vacuum(self, size: int) -> None:
         pass
 
+    @abstractmethod
+    async def storage_use(self) -> Optional[int]:
+        pass
+
+    @abstractmethod
+    async def oldest(self) -> Optional[Key]:
+        pass
+
     async def read_paste_info(self, key: Key) -> Optional[PasteInfo]:
         async with self.connection.execute(
             "select pastes.datetime,pastes.size,pastes.syntax from pastes where pastes.key=? limit 1",
diff --git a/pypaste/server/s3/__init__.py b/pypaste/server/s3/__init__.py
index 6a30358..eeded21 100644
--- a/pypaste/server/s3/__init__.py
+++ b/pypaste/server/s3/__init__.py
@@ -98,33 +98,44 @@ class S3(Storage):
         ) as cursor:
             return await cursor.fetchone() is not None
 
+    async def storage_use(self) -> Optional[int]:
+        async with self.connection.execute(
+            (
+                "select sum(pastes.size) from pastes "
+                "inner join s3 on s3.key=pastes.key"
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [int(use)]:
+                    return use
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
+
+    async def oldest(self) -> Optional[Key]:
+        async with self.connection.execute(
+            (
+                "select pastes.key,pastes.key_length from pastes "
+                "inner join s3 on s3.key=pastes.key "
+                "order by pastes.datetime"
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [bytes(data), int(length)]:
+                    return Key(data, length)
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
+
     async def vacuum(self, max: int) -> None:
         while True:
-            async with self.connection.execute(
-                (
-                    "select sum(pastes.size) from pastes "
-                    "inner join s3 on s3.key "
-                    "where s3.key=pastes.key"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    use = row[0]
-
-            async with self.connection.execute(
-                (
-                    "select pastes.key from pastes "
-                    "inner join s3 on s3.key "
-                    "where s3.key=pastes.key "
-                    "order by pastes.datetime "
-                    "limit 1"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    oldest = row[0]
+            if (use := await self.storage_use()) is None:
+                return
+
+            if (oldest := await self.oldest()) is None:
+                return
 
             if use > max:
                 await self.delete(oldest)
diff --git a/pypaste/server/sqlite/__init__.py b/pypaste/server/sqlite/__init__.py
index be07db6..9ae49ad 100644
--- a/pypaste/server/sqlite/__init__.py
+++ b/pypaste/server/sqlite/__init__.py
@@ -83,33 +83,44 @@ class Sqlite(Storage):
         ) as cursor:
             return await cursor.fetchone() is not None
 
+    async def storage_use(self) -> Optional[int]:
+        async with self.connection.execute(
+            (
+                "select sum(pastes.size) from pastes "
+                "inner join sqlite on sqlite.key=pastes.key"
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [int(use)]:
+                    return use
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
+
+    async def oldest(self) -> Optional[Key]:
+        async with self.connection.execute(
+            (
+                "select pastes.key,pastes.key_length from pastes "
+                "inner join sqlite on sqlite.key=pastes.key "
+                "order by pastes.datetime"
+            )
+        ) as cursor:
+            match await cursor.fetchone():
+                case [data, length]:
+                    return Key(data, length)
+                case None:
+                    return None
+                case _:
+                    raise Exception("unreachable")
+
     async def vacuum(self, max: int) -> None:
         while True:
-            async with self.connection.execute(
-                (
-                    "select sum(pastes.size) from pastes "
-                    "inner join sqlite on sqlite.key "
-                    "where pastes.key=sqlite.key"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    use = row[0]
-
-            async with self.connection.execute(
-                (
-                    "select pastes.key, pastes.key_length from pastes "
-                    "inner join sqlite on sqlite.key "
-                    "where pastes.key=sqlite.key "
-                    "order by pastes.datetime "
-                    "limit 1"
-                )
-            ) as cursor:
-                if (row := await cursor.fetchone()) is None:
-                    return
-                else:
-                    oldest = Key(row[0], row[1])
+            if (use := await self.storage_use()) is None:
+                return
+
+            if (oldest := await self.oldest()) is None:
+                return
 
             if use > max:
                 await self.delete(oldest)
diff --git a/tests/test_storage.py b/tests/test_storage.py
index f2b8912..62dd7b2 100755
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -5,6 +5,7 @@ import os
 import asyncio
 import tempfile
 import aiosqlite
+import zstandard
 from pypaste.server import Paste, Storage, keygen
 from pypaste.server.sqlite import Sqlite
 from pypaste.server.s3 import S3
@@ -17,6 +18,29 @@ def truncate(path: Path) -> None:
         f.truncate(0)
 
 
+async def test_storage_use(storage: Storage) -> None:
+    content = "hello pypaste"
+    length = len(zstandard.compress(content.encode()))
+
+    for x in range(5):
+        paste = Paste(datetime.now(), None, content)
+        key = keygen(6)
+        await storage.insert(paste, key)
+
+    assert (use := await storage.storage_use()) is not None and use == length * 5
+
+
+async def test_oldest(storage: Storage) -> None:
+    pastes = []
+    for x in range(5):
+        paste = Paste(datetime.now(), None, "hello pypaste")
+        key = keygen(6)
+        pastes.append(key)
+        await storage.insert(paste, key)
+
+    assert (oldest := await storage.oldest()) is not None and oldest == pastes[0]
+
+
 async def test_exists_but_not_in_our_table(storage: Storage) -> None:
     key = keygen(6)
 
@@ -125,6 +149,8 @@ async def main() -> int:
         test_delete,
         test_exists,
         test_exists_but_not_in_our_table,
+        test_oldest,
+        test_storage_use,
     ]
 
     await test_sqlite(tests)