1 files changed, 133 insertions, 0 deletions
diff --git a/tests/test_storage.py b/tests/test_storage.py
new file mode 100755
index 0000000..e5f312c
--- /dev/null
+++ b/tests/test_storage.py
@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+import asyncio
+import tempfile
+import aiosqlite
+import string
+import random
+from pypaste.server import Paste, Key, 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
+
+
+def truncate(path: Path) -> None:
+    with open(path, "w") as f:
+        f.truncate(0)
+
+
+def generate_key() -> str:
+    chars = string.ascii_letters
+
+    return "".join(random.choice(chars) for _ in range(10))
+
+
+async def test_exists_but_not_in_our_table(storage: Storage) -> None:
+    key = keygen(6)
+
+    await storage.connection.execute(
+        "insert into pastes values(?, ?, ?, ?, ?)",
+        (key.data, key.length, datetime.now().isoformat(), None, bytes()),
+    )
+
+    assert not await storage.exists(key)
+
+
+async def test_exists(storage: Storage) -> None:
+    dt = datetime.now()
+    key = keygen(6)
+
+    await storage.insert(Paste(dt, "test", "hello world"), key)
+
+    assert await storage.exists(key)
+
+    assert not await storage.exists(keygen(6))
+
+
+async def test_delete(storage: Storage) -> None:
+    dt = datetime.now()
+    key = keygen(6)
+
+    await storage.insert(Paste(dt, "test", "hello world"), key)
+
+    assert await storage.exists(key)
+
+    await storage.delete(key)
+
+    assert not await storage.exists(key)
+
+
+async def test_insert_retrieve(storage: Storage) -> None:
+    dt = datetime.now()
+    key = keygen(6)
+
+    await storage.insert(Paste(dt, "test", "hello world"), key)
+
+    paste = await storage.retrieve(key)
+
+    assert paste is not None
+    assert paste.dt == dt
+    assert paste.syntax == "test"
+    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),
+                )
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(asyncio.run(main()))