diff --git a/app/routers/search.py b/app/routers/search.py index 3f9fca4..d0b0ff9 100644 --- a/app/routers/search.py +++ b/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) diff --git a/pytest.ini b/pytest.ini index b3e7b28..2fc5c9f 100644 --- a/pytest.ini +++ b/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_* diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..f8dcf79 --- /dev/null +++ b/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") diff --git a/tests/test_search_csv.py b/tests/test_search_csv.py new file mode 100644 index 0000000..8fd75ae --- /dev/null +++ b/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]