v1

2026-01-23 20:48:30 +01:00
parent ea3d0164f2
commit 7f0092ff10
9 changed files with 346 additions and 3 deletions
--- a/scripts/storage_quota_test.py
+++ b/scripts/storage_quota_test.py
@@ -0,0 +1,93 @@
+import io
+import os
+import sys
+import tempfile
+
+
+# Ensure repo root is on sys.path when running as a script.
+ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from app import create_app
+from app.extensions import db
+from app.models import Company, Playlist, User
+
+
+def main():
+    # Use a temporary SQLite DB so this doesn't touch your real instance DB.
+    fd, path = tempfile.mkstemp(prefix="rssfeed-test-", suffix=".sqlite")
+    os.close(fd)
+    try:
+        os.environ["DATABASE_URL"] = f"sqlite:///{path}"
+        os.environ["SECRET_KEY"] = "test-secret"
+
+        app = create_app()
+        app.config["TESTING"] = True
+
+        with app.app_context():
+            db.create_all()
+
+            c = Company(name="TestCo")
+            # 1 byte quota so any non-empty upload should be rejected.
+            c.storage_max_bytes = 1
+            db.session.add(c)
+            db.session.commit()
+
+            u = User(username="test@example.com", email="test@example.com", is_admin=False, company_id=c.id)
+            u.set_password("passw0rd123")
+            db.session.add(u)
+            db.session.commit()
+
+            p = Playlist(company_id=c.id, name="Playlist")
+            db.session.add(p)
+            db.session.commit()
+
+            client = app.test_client()
+
+            # Login
+            res = client.post(
+                "/auth/login",
+                data={"email": "test@example.com", "password": "passw0rd123"},
+                follow_redirects=False,
+            )
+            if res.status_code not in (302, 303):
+                raise SystemExit(f"Login failed: {res.status_code}")
+
+            # Try to upload a small image payload.
+            data = {
+                "item_type": "image",
+                "title": "Small",
+                "duration_seconds": "10",
+                "crop_mode": "none",
+                "response": "json",
+                # Not a real image; should fail processing OR quota depending on Pillow.
+                # Use a tiny valid PNG header so Pillow can parse it.
+                "file": (
+                    io.BytesIO(
+                        b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc\xf8\x0f\x00\x01\x01\x01\x00\x18\xdd\x8d\x18\x00\x00\x00\x00IEND\xaeB`\x82"
+                    ),
+                    "tiny.png",
+                ),
+            }
+            url = f"/company/playlists/{p.id}/items"
+            res = client.post(url, data=data, content_type="multipart/form-data")
+
+            if res.status_code != 403:
+                raise SystemExit(f"Expected 403 for quota exceeded, got: {res.status_code} {res.data!r}")
+            js = res.get_json(silent=True) or {}
+            if js.get("ok") is not False:
+                raise SystemExit(f"Unexpected response: {js}")
+            if "Storage limit" not in (js.get("error") or ""):
+                raise SystemExit(f"Unexpected error message: {js}")
+
+            print("OK: storage quota prevents uploads when exceeded")
+    finally:
+        try:
+            os.remove(path)
+        except OSError:
+            pass
+
+
+if __name__ == "__main__":
+    main()