merge main backend app changes

districtr · Oct 9, 2024 · e43cc92 · e43cc92
1 parent 24801ef
commit e43cc92
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 2 deletions.
diff --git a/backend/app/main.py b/backend/app/main.py
@@ -225,7 +225,9 @@ async def get_document(document_id: str, session: Session = Depends(get_session)
 async def get_total_population(
     document_id: str, session: Session = Depends(get_session)
 ):
-    stmt = text("SELECT * from get_total_population(:document_id)")
+    stmt = text(
+        "SELECT * from get_total_population(:document_id) WHERE zone IS NOT NULL"
+    )
     try:
         result = session.execute(stmt, {"document_id": document_id})
         return [

diff --git a/backend/app/models.py b/backend/app/models.py
@@ -132,7 +132,7 @@ class AssignmentsBase(SQLModel):
     metadata = MetaData(schema=DOCUMENT_SCHEMA)
     document_id: str = Field(sa_column=Column(UUIDType, primary_key=True))
     geo_id: str = Field(primary_key=True)
-    zone: int
+    zone: int | None
 
 
 class Assignments(AssignmentsBase, table=True):

diff --git a/backend/tests/test_main.py b/backend/tests/test_main.py
@@ -322,6 +322,21 @@ def test_patch_assignments(client, document_id):
     assert response.json() == {"assignments_upserted": 3}
 
 
+def test_patch_assignments_nulls(client, document_id):
+    response = client.patch(
+        "/api/update_assignments",
+        json={
+            "assignments": [
+                {"document_id": document_id, "geo_id": "202090416004010", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090416003004", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090434001003", "zone": None},
+            ]
+        },
+    )
+    assert response.status_code == 200
+    assert response.json() == {"assignments_upserted": 3}
+
+
 def test_patch_assignments_twice(client, document_id):
     response = client.patch(
         "/api/update_assignments",
@@ -359,6 +374,28 @@ def test_patch_assignments_twice(client, document_id):
     assert data[1]["geo_id"] == "202090434001003"
 
 
+def test_get_document_population_totals_null_assignments(
+    client, document_id, ks_demo_view_census_blocks
+):
+    response = client.patch(
+        "/api/update_assignments",
+        json={
+            "assignments": [
+                {"document_id": document_id, "geo_id": "202090416004010", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090416003004", "zone": 1},
+                {"document_id": document_id, "geo_id": "202090434001003", "zone": None},
+            ]
+        },
+    )
+    assert response.status_code == 200
+    assert response.json() == {"assignments_upserted": 3}
+    doc_uuid = str(uuid.UUID(document_id))
+    result = client.get(f"/api/document/{doc_uuid}/total_pop")
+    assert result.status_code == 200
+    data = result.json()
+    assert data == [{"zone": 1, "total_pop": 67}]
+
+
 def test_get_document_population_totals(
     client, assignments_document_id, ks_demo_view_census_blocks
 ):