1 files changed, 55 insertions, 50 deletions
diff --git a/tests/test_storage.py b/tests/test_storage.py
index dde2f40..029e413 100755
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -5,12 +5,14 @@ import os
 import asyncio
 import tempfile
 import aiosqlite
+import zstandard
+import string
 from pypaste.server import Paste, Storage, keygen
 from pypaste.server.sqlite import Sqlite
 from pypaste.server.s3 import S3
 from datetime import datetime
 from pathlib import Path
-from typing import List
+from secrets import choice
 
 
 def truncate(path: Path) -> None:
@@ -18,6 +20,23 @@ def truncate(path: Path) -> None:
         f.truncate(0)
 
 
+def random_string(length: int) -> str:
+    return "".join(choice(string.ascii_letters) for _ in range(length))
+
+
+async def test_vacuum(storage: Storage):
+    # generate random string to avoid it getting compressed
+    content = random_string(32)
+    size_compressed = len(zstandard.compress(content.encode()))
+
+    for x in range(1024 // size_compressed + 1):
+        await storage.insert(Paste(datetime.now(), None, content), keygen(6))
+
+    await storage.vacuum(256)
+
+    assert (use := await storage.get_used()) is not None and use <= 256
+
+
 async def test_exists_but_not_in_our_table(storage: Storage) -> None:
     key = keygen(6)
 
@@ -67,57 +86,43 @@ async def test_insert_retrieve(storage: Storage) -> None:
     assert paste.text == "hello world"
 
 
-async def main() -> int:
-    stores: List[Storage] = []
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        f = Path(tmpdir) / "database"
-        truncate(f)
-        async with aiosqlite.connect(f) as connection:
-            await connection.execute(
-                (
-                    "create table pastes("
-                    "key blob,"
-                    "key_length int,"
-                    "datetime text,"
-                    "size int,"
-                    "syntax text"
-                    ")"
-                )
-            )
-
-            sqlite_storage = Sqlite(connection)
-            await sqlite_storage.setup()
-            stores.append(sqlite_storage)
-
-            try:
-                os.environ["PYPASTE_TEST_S3"]
-                test_s3 = True
-            except KeyError:
-                test_s3 = False
-
-            if test_s3:
-                s3_storage = S3(
-                    connection,
-                    os.environ["PYPASTE_TEST_ENDPOINT"],
-                    os.environ["PYPASTE_TEST_REGION"],
-                    os.environ["PYPASTE_TEST_BUCKET"],
-                    os.environ["PYPASTE_TEST_ACCESS_KEY"],
-                    os.environ["PYPASTE_TEST_SECRET_KEY"],
-                )
-                await s3_storage.setup()
-                stores.append(s3_storage)
-
-            for store in stores:
-                await asyncio.gather(
-                    test_insert_retrieve(store),
-                    test_insert_retrieve(store),
-                    test_delete(store),
-                    test_delete(store),
-                    test_exists_but_not_in_our_table(store),
-                    test_exists_but_not_in_our_table(store),
+async def test_sqlite(tests):
+    for test in tests:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            f = Path(tmpdir) / "pastes.sqlite"
+            truncate(f)
+
+            async with aiosqlite.connect(f) as connection:
+                await connection.execute(
+                    (
+                        "create table pastes("
+                        "key blob,"
+                        "key_length int,"
+                        "datetime text,"
+                        "size int,"
+                        "syntax text"
+                        ")"
+                    )
                 )
 
+                storage = Sqlite(connection)
+
+                await storage.setup()
+
+                await test(storage)
+
+
+async def main() -> int:
+    tests = [
+        test_insert_retrieve,
+        test_delete,
+        test_exists,
+        test_exists_but_not_in_our_table,
+        test_vacuum,
+    ]
+
+    await test_sqlite(tests)
+
     return 0