add aurweb.models.ban.Ban ORM mapping

Signed-off-by: Kevin Morris <kevr@0cost.org>

aurweb/models/ban.py

@@ -0,0 +1,19 @@
+from fastapi import Request
+from sqlalchemy.orm import mapper
+
+from aurweb.schema import Bans
+
+
+class Ban:
+    def __init__(self, **kwargs):
+        self.IPAddress = kwargs.get("IPAddress")
+        self.BanTS = kwargs.get("BanTS")
+
+
+def is_banned(request: Request):
+    from aurweb.db import session
+    ip = request.client.host
+    return session.query(Ban).filter(Ban.IPAddress == ip).first() is not None
+
+
+mapper(Ban, Bans)

test/test_ban.py

@@ -0,0 +1,63 @@
+import warnings
+
+from datetime import datetime, timedelta
+
+import pytest
+
+from sqlalchemy import exc as sa_exc
+
+from aurweb.models.ban import Ban, is_banned
+from aurweb.testing import setup_test_db
+from aurweb.testing.requests import Request
+
+ban = None
+
+request = Request()
+
+
+@pytest.fixture(autouse=True)
+def setup():
+    from aurweb.db import session
+
+    global ban
+
+    setup_test_db("Bans")
+
+    ban = Ban(IPAddress="127.0.0.1",
+              BanTS=datetime.utcnow() + timedelta(seconds=30))
+    session.add(ban)
+    session.commit()
+
+
+def test_ban():
+    assert ban.IPAddress == "127.0.0.1"
+    assert bool(ban.BanTS)
+
+
+def test_invalid_ban():
+    from aurweb.db import session
+
+    with pytest.raises(sa_exc.IntegrityError,
+                       match="NOT NULL constraint failed: Bans.IPAddress"):
+        bad_ban = Ban(BanTS=datetime.utcnow())
+        session.add(bad_ban)
+
+        # We're adding a ban with no primary key; this causes an
+        # SQLAlchemy warnings when committing to the DB.
+        # Ignore them.
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", sa_exc.SAWarning)
+            session.commit()
+
+    # Since we got a transaction failure, we need to rollback.
+    session.rollback()
+
+
+def test_banned():
+    request.client.host = "127.0.0.1"
+    assert is_banned(request)
+
+
+def test_not_banned():
+    request.client.host = "192.168.0.1"
+    assert not is_banned(request)