Fix CSV search download

app/routers/search.py

@@ -11,7 +11,7 @@
 from app.db.database import Database, get_db
 from app.db.queries import get_ec_suggestions, get_filtered_data, get_total_count, get_typeahead_suggestions
 from app.models.query_params import CLEANECLookupQueryParams, CLEANSearchQueryParams, CLEANTypeaheadQueryParams, ResponseFormat
-from app.models.clean_data import CLEANDataBase, CLEANECLookupResponse, CLEANECLookupMatch, CLEANSearchResponse, CLEANTypeaheadResponse, CurationStatusOption, CLEANCurationStatusResponse
+from app.models.clean_data import CLEANColumn, CLEANDataBase, CLEANECLookupResponse, CLEANECLookupMatch, CLEANSearchResponse, CLEANTypeaheadResponse, CurationStatusOption, CLEANCurationStatusResponse
 
 router = APIRouter(tags=["Search"])
 
@@ -219,21 +219,17 @@ async def get_data(
             # Create CSV response
             output = StringIO()
 
-            # Determine columns to include in CSV
-            if params.columns:
-                fieldnames = params.columns
-            elif data and len(data) > 0:
-                fieldnames = list(data[0].keys())
-            else:
-                fieldnames = []
-
             # Write CSV
             writer = csv.DictWriter(
-                output, fieldnames=fieldnames, extrasaction="ignore"
+                output, fieldnames=[c.value for c in CLEANColumn], extrasaction="ignore"
             )
             writer.writeheader()
             for row in data:
-                writer.writerow(row)
+                # Join array columns into one ';'-separated cell
+                writer.writerow({
+                    k: ";".join(str(v) for v in value) if isinstance(value, list) else value
+                    for k, value in row.items()
+                })
 
             # Return streaming response
             output.seek(0)

pytest.ini

@@ -2,6 +2,7 @@
 asyncio_mode = auto
 asyncio_default_fixture_loop_scope = function
 testpaths = tests
+pythonpath = .
 python_files = test_*.py
 python_classes = Test*
 python_functions = test_*

tests/conftest.py

@@ -0,0 +1,5 @@
+import os
+
+# Settings require DB credentials at import time; tests mock the DB calls.
+for key in ("CLEAN_DB_HOST", "CLEAN_DB_USER", "CLEAN_DB_PASSWORD", "CLEAN_DB_NAME"):
+    os.environ.setdefault(key, "test")

tests/test_search_csv.py

@@ -0,0 +1,56 @@
+import csv
+from io import StringIO
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from fastapi.testclient import TestClient
+
+import app.routers.search as search
+from app.db.database import get_db
+from app.main import app
+from app.models.clean_data import CLEANColumn
+
+HEADER = [c.value for c in CLEANColumn]
+
+
+@pytest.fixture
+def client():
+    app.dependency_overrides[get_db] = lambda: None
+    yield TestClient(app)
+    app.dependency_overrides.clear()
+
+
+def get_csv(client, rows):
+    with (
+        patch.object(search, "get_total_count", AsyncMock(return_value=len(rows))),
+        patch.object(search, "get_filtered_data", AsyncMock(return_value=rows)),
+    ):
+        response = client.get("/api/v1/search?accession=P51451&format=csv&limit=10")
+    assert response.status_code == 200, response.text
+    return list(csv.reader(StringIO(response.text)))
+
+
+def make_row(**overrides):
+    row = {c.value: None for c in CLEANColumn}
+    row.update(accession="P51451", protein_name='Kinase, "BLK"')
+    row.update(overrides)
+    return row
+
+
+def test_csv_joins_array_columns(client):
+    rows = get_csv(client, [make_row(
+        clean_ec_number_array=["2.7.10.2", "2.7.11.1"],
+        clean_ec_confidence_array=[0.98, 0.12],
+        annot_ec_number_array=None,
+    )])
+
+    assert rows[0] == HEADER
+    record = dict(zip(HEADER, rows[1]))
+    assert record["protein_name"] == 'Kinase, "BLK"'
+    assert record["clean_ec_number_array"] == "2.7.10.2;2.7.11.1"
+    assert record["clean_ec_confidence_array"] == "0.98;0.12"
+    assert record["annot_ec_number_array"] == ""
+
+
+def test_csv_has_header_when_no_rows_match(client):
+    assert get_csv(client, []) == [HEADER]