From 2da33e86fceb61678f0b7370f85064bd381a377f Mon Sep 17 00:00:00 2001
From: Alyssa Dai <alyssa.ydai@gmail.com>
Date: Thu, 24 Oct 2024 15:37:07 -0400
Subject: [PATCH] [FIX] Allow only `"true"` or `None` for `is_control` query
 parameter (#129)

* change is_control to str in QueryModel & add basic validation

* update type hints in CRUD func

* add simple tests of query model field validation
---
 app/api/crud.py     |  4 ++--
 app/api/models.py   | 19 +++++++++++++++++--
 tests/test_query.py | 21 +++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/app/api/crud.py b/app/api/crud.py
index 0920a4c..4c38dc3 100644
--- a/app/api/crud.py
+++ b/app/api/crud.py
@@ -37,7 +37,7 @@ async def get(
     max_age: float,
     sex: str,
     diagnosis: str,
-    is_control: bool,
+    is_control: str,
     min_num_imaging_sessions: int,
     min_num_phenotypic_sessions: int,
     assessment: str,
@@ -59,7 +59,7 @@ async def get(
         Sex of subject.
     diagnosis : str
         Subject diagnosis.
-    is_control : bool
+    is_control : str
         Whether or not subject is a control.
     min_num_imaging_sessions : int
         Subject minimum number of imaging sessions.
diff --git a/app/api/models.py b/app/api/models.py
index 81cb48c..e0c04a4 100644
--- a/app/api/models.py
+++ b/app/api/models.py
@@ -4,7 +4,8 @@
 from typing import Optional, Union
 
 from fastapi import Query
-from pydantic import BaseModel, Field
+from fastapi.exceptions import HTTPException
+from pydantic import BaseModel, Field, validator
 
 CONTROLLED_TERM_REGEX = r"^[a-zA-Z]+[:]\S+$"
 
@@ -16,7 +17,7 @@ class QueryModel(BaseModel):
     max_age: float = None
     sex: str = None
     diagnosis: str = None
-    is_control: bool = None
+    is_control: str = None
     min_num_imaging_sessions: int = None
     min_num_phenotypic_sessions: int = None
     assessment: str = None
@@ -27,6 +28,20 @@ class QueryModel(BaseModel):
     # syntax from https://github.com/tiangolo/fastapi/issues/4445#issuecomment-1117632409
     node_url: list[str] | None = Field(Query(default=[]))
 
+    @validator("is_control")
+    def check_allowed_iscontrol_values(cls, v):
+        """Raise a validation error if the value of 'is_control' is not 'true' (case-insensitive) or None."""
+        if v is not None:
+            # Ensure that the allowed value is case-insensitive
+            if v.lower() != "true":
+                raise HTTPException(
+                    status_code=422,
+                    detail="'is_control' must be either set to 'true' or omitted from the query",
+                )
+            # Keep it a str because that's what the n-API expects
+            return v
+        return None
+
 
 class CohortQueryResponse(BaseModel):
     """Data model for query results for one matching dataset (i.e., a cohort)."""
diff --git a/tests/test_query.py b/tests/test_query.py
index 44f9dd9..5891005 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -4,6 +4,9 @@
 import httpx
 import pytest
 from fastapi import status
+from fastapi.exceptions import HTTPException
+
+from app.api.models import QueryModel
 
 
 @pytest.fixture()
@@ -241,3 +244,21 @@ async def mock_httpx_get(self, **kwargs):
     response = test_app.get("/query")
 
     assert response.status_code == status.HTTP_200_OK
+
+
+@pytest.mark.parametrize("valid_iscontrol", ["true", "True", "TRUE"])
+def test_valid_iscontrol_values_unchanged(valid_iscontrol):
+    """Test that valid values for the 'is_control' field are accepted without modification or errors."""
+    example_fquery = QueryModel(is_control=valid_iscontrol)
+    assert example_fquery.is_control == valid_iscontrol
+
+
+@pytest.mark.parametrize("invalid_iscontrol", ["false", 52])
+def test_invalid_iscontrol_value_raises_error(invalid_iscontrol):
+    """Test that invalid values for the 'is_control' field fail model validation and raise an informative HTTPException."""
+    with pytest.raises(HTTPException) as exc:
+        QueryModel(is_control=invalid_iscontrol)
+
+    assert "'is_control' must be either set to 'true' or omitted" in str(
+        exc.value
+    )