Add test for latest endpoint

cid/main.py

@@ -1,9 +1,9 @@
 import logging
-from typing import Any, Dict
+from typing import Any, Dict, Generator
 
-from fastapi import FastAPI, Request
+from fastapi import Depends, FastAPI
 from sqlalchemy import create_engine, desc
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import Session, sessionmaker
 
 from cid.config import DATABASE_URL
 from cid.models import AwsImage, AzureImage, GoogleImage
@@ -15,13 +15,20 @@
 app = FastAPI()
 
 
+def get_db() -> Generator:
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+
 @app.get("/")
 def read_root() -> dict:
     return {"Hello": "World"}
 
 
-def latest_aws_image(request: Request) -> Dict[str, Any]:
-    db = SessionLocal()
+def latest_aws_image(db: Session) -> Dict[str, Any]:
     regions = db.query(AwsImage.region).distinct().order_by(AwsImage.region).all()
     latest_image = (
         db.query(AwsImage).order_by(desc(AwsImage.version), desc(AwsImage.date)).first()
@@ -48,7 +55,7 @@ def latest_aws_image(request: Request) -> Dict[str, Any]:
     }
 
 
-def latest_azure_image(request: Request) -> Dict[str, Any]:
+def latest_azure_image(db: Session) -> Dict[str, Any]:
     db = SessionLocal()
     latest_image = db.query(AzureImage).order_by(desc(AzureImage.version)).first()
 
@@ -63,7 +70,7 @@ def latest_azure_image(request: Request) -> Dict[str, Any]:
     }
 
 
-def latest_google_image(request: Request) -> Dict[str, Any]:
+def latest_google_image(db: Session) -> Dict[str, Any]:
     db = SessionLocal()
     latest_image = (
         db.query(GoogleImage)
@@ -83,9 +90,9 @@ def latest_google_image(request: Request) -> Dict[str, Any]:
 
 
 @app.get("/latest")
-def latest(request: Request) -> Dict[str, Any]:
+def latest(db: Session = Depends(get_db)) -> Dict[str, Any]:  # noqa: B008
     return {
-        "latest_aws_image": latest_aws_image(request),
-        "latest_azure_image": latest_azure_image(request),
-        "latest_google_image": latest_google_image(request),
+        "latest_aws_image": latest_aws_image(db),
+        "latest_azure_image": latest_azure_image(db),
+        "latest_google_image": latest_google_image(db),
     }

pyproject.toml

@@ -23,6 +23,7 @@ pre-commit = "^3.7.1"
 tox = "^4.15.0"
 pytest-randomly = "^3.15.0"
 pytest-sugar = "^1.0.0"
+sqlalchemy-utils = "^0.41.2"
 
 [build-system]
 requires = ["poetry-core"]

tests/test_main.py

@@ -1,8 +1,30 @@
 """Tests for the main module."""
 
+from datetime import datetime
+
 from fastapi.testclient import TestClient
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy_utils import drop_database
+
+from cid.main import app, get_db, latest_aws_image
+from cid.models import AwsImage
+
+# Create an in-memory SQLite database for testing
+DATABASE_URL = "sqlite:///./test.db"
+engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
+TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+
 
-from cid.main import app
+def override_get_db():
+    try:
+        db = TestingSessionLocal()
+        yield db
+    finally:
+        db.close()
+
+
+app.dependency_overrides[get_db] = override_get_db
 
 # Create a TestClient instance to test the FastAPI app
 # https://fastapi.tiangolo.com/tutorial/testing/
@@ -14,3 +36,38 @@ def test_read_root():
     response = client.get("/")
     assert response.status_code == 200
     assert response.json() == {"Hello": "World"}
+
+
+def teardown():
+    AwsImage.metadata.drop_all(bind=engine)
+    TestingSessionLocal.remove()
+    engine.dispose()
+
+
+def test_latest_aws_image():
+    AwsImage.metadata.create_all(bind=engine)
+
+    db = TestingSessionLocal()
+    aws_image = AwsImage(
+        id="ami-12345678",
+        name="test_image",
+        version="1.0",
+        # SQLite only supports datetime objects
+        date=datetime.strptime("2022-01-01", "%Y-%m-%d").date(),
+        region="us-west-1",
+        imageId="ami-12345678",
+    )
+    db.add(aws_image)
+    db.commit()
+
+    response = latest_aws_image(db)
+
+    assert response == {
+        "name": "test_image",
+        "version": "1.0",
+        "date": datetime(2022, 1, 1, 0, 0),
+        "amis": {"us-west-1": "ami-12345678"},
+    }
+
+    drop_database(engine.url)
+    app.dependency_overrides.pop(get_db)