Add type to get_or_404, update tests

Author: pamelafox

.gitignore

@@ -17,3 +17,4 @@ env/
 venv/
 .venv/
 .pdm.toml
+.pdm-python

src/flask_sqlalchemy/extension.py

@@ -23,6 +23,8 @@
 from .session import Session
 from .table import _Table
 
+_O = t.TypeVar("_O", bound=object)
+
 
 class SQLAlchemy:
     """Integrates SQLAlchemy with Flask. This handles setting up one or more engines,
@@ -730,8 +732,8 @@ def get_binds(self) -> dict[sa.Table, sa.engine.Engine]:
         }
 
     def get_or_404(
-        self, entity: type[t.Any], ident: t.Any, *, description: str | None = None
-    ) -> t.Any:
+        self, entity: type[_O], ident: t.Any, *, description: str | None = None
+    ) -> t.Optional[_O]:
         """Like :meth:`session.get() <sqlalchemy.orm.Session.get>` but aborts with a
         ``404 Not Found`` error instead of returning ``None``.
 

tests/conftest.py

@@ -31,7 +31,7 @@ def db(app: Flask) -> SQLAlchemy:
 
 
 @pytest.fixture
-def Todo(app: Flask, db: SQLAlchemy) -> t.Any:
+def Todo(app: Flask, db: SQLAlchemy) -> t.Generator[t.Any, None, None]:
     class Todo(db.Model):
         id = sa.Column(sa.Integer, primary_key=True)
         title = sa.Column(sa.String)

tests/test_view_query.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import typing as t
+
+import pytest
+import sqlalchemy as sa
+from flask import Flask
+from werkzeug.exceptions import NotFound
+
+from flask_sqlalchemy import SQLAlchemy
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_view_get_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    item = Todo()
+    db.session.add(item)
+    db.session.commit()
+    assert db.get_or_404(Todo, 1) is item
+    with pytest.raises(NotFound):
+        assert db.get_or_404(Todo, 2)
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_first_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add(Todo(title="a"))
+    db.session.commit()
+    result = db.first_or_404(db.select(Todo).filter_by(title="a"))
+    assert result.title == "a"
+
+    with pytest.raises(NotFound):
+        db.first_or_404(db.select(Todo).filter_by(title="b"))
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_view_one_or_404(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add(Todo(title="a"))
+    db.session.add(Todo(title="b"))
+    db.session.add(Todo(title="b"))
+    db.session.commit()
+    result = db.one_or_404(db.select(Todo).filter_by(title="a"))
+    assert result.title == "a"
+
+    with pytest.raises(NotFound):
+        # MultipleResultsFound
+        db.one_or_404(db.select(Todo).filter_by(title="b"))
+
+    with pytest.raises(NotFound):
+        # NoResultFound
+        db.one_or_404(db.select(Todo).filter_by(title="c"))
+
+
+@pytest.mark.usefixtures("app_ctx")
+def test_paginate(db: SQLAlchemy, Todo: t.Any) -> None:
+    db.session.add_all(Todo() for _ in range(150))
+    db.session.commit()
+    p = db.paginate(db.select(Todo))
+    assert p.total == 150
+    assert len(p.items) == 20
+    p2 = p.next()
+    assert p2.page == 2
+    assert p2.total == 150
+
+
+# This test creates its own inline model so that it can use that as the type
+@pytest.mark.usefixtures("app_ctx")
+def test_view_get_or_404_typed(db: SQLAlchemy, app: Flask) -> None:
+    class Quiz(db.Model):
+        id = sa.Column(sa.Integer, primary_key=True)
+        topic = sa.Column(sa.String)
+
+    db.create_all()
+
+    item: Quiz = Quiz()
+    db.session.add(item)
+    db.session.commit()
+    result: Quiz = db.get_or_404(Quiz, 1)
+    assert result is item
+    with pytest.raises(NotFound):
+        assert db.get_or_404(Quiz, 2)
+    db.drop_all()