update existing integration tests

Author: jakewheeler

refiner/tests/integration/test_export.py

@@ -1,6 +1,7 @@
 import csv
 import re
-from io import StringIO
+import zipfile
+from io import BytesIO, StringIO
 
 import pytest
 from fastapi import status
@@ -10,6 +11,15 @@
 from app.db.configurations.model import DbConfigurationCustomCode
 
 
+def get_csv_from_zip(content: bytes, filename_pattern: str) -> str:
+    """Extract a CSV file from a zip response by filename pattern."""
+    with zipfile.ZipFile(BytesIO(content)) as zf:
+        for name in zf.namelist():
+            if re.search(filename_pattern, name):
+                return zf.read(name).decode("utf-8")
+    raise FileNotFoundError(f"No file matching {filename_pattern!r} found in zip")
+
+
 @pytest.mark.integration
 @pytest.mark.asyncio
 class TestConfigurationExport:
@@ -21,11 +31,11 @@ async def test_export_returns_404_for_unknown_id(self, setup, authed_client):
         response = await authed_client.get(f"/api/v1/configurations/{dummy_id}/export")
         assert response.status_code == status.HTTP_404_NOT_FOUND
 
-    async def test_export_returns_csv_with_correct_headers(
+    async def test_export_returns_zip_with_correct_headers(
         self, setup, authed_client, get_condition_id, create_config
     ):
         """
-        CSV should be returned in correct form when given a valid config.
+        Zip should be returned in correct form when given a valid config.
         """
         condition_id = await get_condition_id("Colorado tick fever")
         config = await create_config(condition_id)
@@ -34,14 +44,33 @@ async def test_export_returns_csv_with_correct_headers(
         )
 
         assert response.status_code == status.HTTP_200_OK
-        assert response.headers["content-type"].startswith("text/csv")
+        assert response.headers["content-type"].startswith("application/zip")
 
         cd_header = response.headers.get("content-disposition", "")
         assert re.search(
-            r'filename=".+_Code_Export_\d{6}_\d{2}_\d{2}_\d{2}\.csv"',
+            r'filename=".+_Configuration_Export_\d{6}_\d{2}_\d{2}_\d{2}\.zip"',
             cd_header,
         ), f"Unexpected Content-Disposition: {cd_header!r}"
 
+    async def test_export_zip_contains_expected_files(
+        self, setup, authed_client, get_condition_id, create_config
+    ):
+        """
+        Zip should contain a codes CSV and a sections CSV.
+        """
+        condition_id = await get_condition_id("Colorado tick fever")
+        config = await create_config(condition_id)
+        response = await authed_client.get(
+            f"/api/v1/configurations/{config['id']}/export"
+        )
+
+        assert response.status_code == status.HTTP_200_OK
+
+        with zipfile.ZipFile(BytesIO(response.content)) as zf:
+            names = zf.namelist()
+            assert any(re.search(r"Code_Export.+\.csv", n) for n in names)
+            assert any(re.search(r"Section_Export.+\.csv", n) for n in names)
+
     async def test_export_includes_all_codes_from_multiple_codesets(
         self,
         setup,
@@ -72,7 +101,7 @@ async def test_export_includes_all_codes_from_multiple_codesets(
         response = await authed_client.get(f"/api/v1/configurations/{config_id}/export")
         assert response.status_code == status.HTTP_200_OK
 
-        content = response.text
+        content = get_csv_from_zip(response.content, r"Code_Export")
         lines = [line for line in content.splitlines() if line.strip()]
         assert len(lines) == expected_code_rows + 1  # header + codes
 
@@ -87,7 +116,6 @@ async def test_export_csv_conditions_column_correct(
         """
         CSV "Condition" column should only contain the associated condition names.
         """
-
         amebiasis_id = await get_condition_id("Amebiasis")
         config = await create_config(amebiasis_id)
         config_id = config["id"]
@@ -98,9 +126,9 @@ async def test_export_csv_conditions_column_correct(
         response = await authed_client.get(f"/api/v1/configurations/{config_id}/export")
         assert response.status_code == status.HTTP_200_OK
 
-        reader = csv.DictReader(StringIO(response.text))
-        column_name = "Condition"
-        conditions_in_csv = {row[column_name] for row in reader if row[column_name]}
+        content = get_csv_from_zip(response.content, r"Code_Export")
+        reader = csv.DictReader(StringIO(content))
+        conditions_in_csv = {row["Condition"] for row in reader if row["Condition"]}
 
         assert conditions_in_csv == {"Amebiasis", "Byssinosis"}
 
@@ -132,17 +160,18 @@ async def test_export_csv_code_systems_valid(
         response = await authed_client.get(f"/api/v1/configurations/{config_id}/export")
         assert response.status_code == status.HTTP_200_OK
 
-        reader = csv.DictReader(StringIO(response.text))
-        column_name = "Code System"
-        code_systems_in_csv = {row[column_name] for row in reader if row[column_name]}
+        content = get_csv_from_zip(response.content, r"Code_Export")
+        reader = csv.DictReader(StringIO(content))
+        code_systems_in_csv = {
+            row["Code System"] for row in reader if row["Code System"]
+        }
 
         code_systems = await get_all_code_systems_db(db=db_pool)
         expected_systems = {cs.display_name for cs in code_systems.values()}
 
-        # csv systems must be a subset of the expected systems set
         assert code_systems_in_csv <= expected_systems
 
-    async def test_export_custom_codes_have_blank_code_system(
+    async def test_export_custom_codes_have_blank_condition(
         self,
         setup,
         authed_client,
@@ -169,18 +198,19 @@ async def test_export_custom_codes_have_blank_code_system(
         response = await authed_client.get(f"/api/v1/configurations/{config_id}/export")
         assert response.status_code == status.HTTP_200_OK
 
-        reader = csv.DictReader(StringIO(response.text))
+        content = get_csv_from_zip(response.content, r"Code_Export")
+        reader = csv.DictReader(StringIO(content))
         for row in reader:
             if row["Code Type"] == "Custom code":
                 assert row["Condition"] == "", (
-                    f"Expected blank Code System for custom code, got {row['Code System']!r}"
+                    f"Expected blank Condition for custom code, got {row['Condition']!r}"
                 )
 
     async def test_export_csv_body_is_non_empty(
         self, setup, authed_client, get_condition_id, create_config
     ):
         """
-        CSV should contain at least a header row.
+        Codes CSV should contain at least a header row.
         """
         condition_id = await get_condition_id("Cholera")
         config = await create_config(condition_id)
@@ -189,6 +219,6 @@ async def test_export_csv_body_is_non_empty(
         )
 
         assert response.status_code == status.HTTP_200_OK
-        content = response.text
+        content = get_csv_from_zip(response.content, r"Code_Export")
         lines = [line for line in content.splitlines() if line.strip()]
         assert len(lines) >= 1, "Expected at least a CSV header row in the response"