"""Tests for cohort builder query_builder module.
Tests cover:
- Empty criteria (no WHERE clause)
- Each criterion in isolation
- Combined criteria
- Validation errors (age boundaries, invalid gender, ICD injection)
- All valid SQL passes is_safe_query()
"""
import pytest
from m4.apps.cohort_builder.query_builder import (
MAX_AGE,
MIN_AGE,
QueryCohortInput,
build_cohort_count_sql,
build_cohort_demographics_sql,
build_gender_distribution_sql,
)
from m4.core.validation import is_safe_query
class TestEmptyCriteria:
"""Test queries with no filtering criteria."""
def test_count_sql_no_where_clause(self):
"""Empty criteria should produce SQL without WHERE clause."""
criteria = QueryCohortInput()
sql = build_cohort_count_sql(criteria)
assert "WHERE" not in sql
assert "mimiciv_hosp.patients" in sql
assert "mimiciv_hosp.admissions" in sql
assert "patient_count" in sql
assert "admission_count" in sql
def test_demographics_sql_no_where_clause(self):
"""Empty criteria should produce demographics SQL without WHERE clause."""
criteria = QueryCohortInput()
sql = build_cohort_demographics_sql(criteria)
assert "WHERE" not in sql
assert "age_bucket" in sql
def test_gender_sql_no_where_clause(self):
"""Empty criteria should produce gender SQL without WHERE clause."""
criteria = QueryCohortInput()
sql = build_gender_distribution_sql(criteria)
assert "WHERE" not in sql
assert "p.gender" in sql
def test_empty_criteria_passes_validation(self):
"""Empty criteria SQL should pass safety validation."""
criteria = QueryCohortInput()
for builder in [
build_cohort_count_sql,
build_cohort_demographics_sql,
build_gender_distribution_sql,
]:
sql = builder(criteria)
safe, msg = is_safe_query(sql)
assert safe, f"SQL failed validation: {msg}"
class TestAgeCriteria:
"""Test age filtering criteria."""
def test_age_min_only(self):
"""age_min should add >= condition."""
criteria = QueryCohortInput(age_min=18)
sql = build_cohort_count_sql(criteria)
assert "WHERE" in sql
assert "p.anchor_age >= 18" in sql
def test_age_max_only(self):
"""age_max should add <= condition."""
criteria = QueryCohortInput(age_max=65)
sql = build_cohort_count_sql(criteria)
assert "WHERE" in sql
assert "p.anchor_age <= 65" in sql
def test_age_range(self):
"""Both age_min and age_max should create range."""
criteria = QueryCohortInput(age_min=18, age_max=65)
sql = build_cohort_count_sql(criteria)
assert "p.anchor_age >= 18" in sql
assert "p.anchor_age <= 65" in sql
assert " AND " in sql
def test_age_at_boundary_min(self):
"""Age at MIN_AGE boundary should be valid."""
criteria = QueryCohortInput(age_min=MIN_AGE)
sql = build_cohort_count_sql(criteria)
assert f"p.anchor_age >= {MIN_AGE}" in sql
def test_age_at_boundary_max(self):
"""Age at MAX_AGE boundary should be valid."""
criteria = QueryCohortInput(age_max=MAX_AGE)
sql = build_cohort_count_sql(criteria)
assert f"p.anchor_age <= {MAX_AGE}" in sql
class TestAgeValidation:
"""Test age validation errors."""
def test_age_min_negative(self):
"""Negative age_min should raise ValueError."""
criteria = QueryCohortInput(age_min=-1)
with pytest.raises(ValueError, match="age_min must be between"):
build_cohort_count_sql(criteria)
def test_age_max_negative(self):
"""Negative age_max should raise ValueError."""
criteria = QueryCohortInput(age_max=-1)
with pytest.raises(ValueError, match="age_max must be between"):
build_cohort_count_sql(criteria)
def test_age_min_too_high(self):
"""age_min above MAX_AGE should raise ValueError."""
criteria = QueryCohortInput(age_min=MAX_AGE + 1)
with pytest.raises(ValueError, match="age_min must be between"):
build_cohort_count_sql(criteria)
def test_age_max_too_high(self):
"""age_max above MAX_AGE should raise ValueError."""
criteria = QueryCohortInput(age_max=MAX_AGE + 1)
with pytest.raises(ValueError, match="age_max must be between"):
build_cohort_count_sql(criteria)
def test_age_min_greater_than_max(self):
"""age_min > age_max should raise ValueError."""
criteria = QueryCohortInput(age_min=65, age_max=18)
with pytest.raises(ValueError, match="cannot be greater than"):
build_cohort_count_sql(criteria)
class TestGenderCriteria:
"""Test gender filtering criteria."""
def test_gender_male(self):
"""Gender 'M' should add proper condition."""
criteria = QueryCohortInput(gender="M")
sql = build_cohort_count_sql(criteria)
assert "p.gender = 'M'" in sql
def test_gender_female(self):
"""Gender 'F' should add proper condition."""
criteria = QueryCohortInput(gender="F")
sql = build_cohort_count_sql(criteria)
assert "p.gender = 'F'" in sql
class TestGenderValidation:
"""Test gender validation errors."""
def test_invalid_gender_x(self):
"""Invalid gender 'X' should raise ValueError."""
criteria = QueryCohortInput(gender="X")
with pytest.raises(ValueError, match="gender must be one of"):
build_cohort_count_sql(criteria)
def test_invalid_gender_lowercase(self):
"""Lowercase gender 'm' should raise ValueError."""
criteria = QueryCohortInput(gender="m")
with pytest.raises(ValueError, match="gender must be one of"):
build_cohort_count_sql(criteria)
def test_invalid_gender_empty(self):
"""Empty gender string should raise ValueError."""
criteria = QueryCohortInput(gender="")
with pytest.raises(ValueError, match="gender must be one of"):
build_cohort_count_sql(criteria)
class TestIcdCodesCriteria:
"""Test ICD codes filtering criteria."""
def test_single_icd_code(self):
"""Single ICD code should create EXISTS subquery."""
criteria = QueryCohortInput(icd_codes=["I10"])
sql = build_cohort_count_sql(criteria)
assert "EXISTS" in sql
assert "mimiciv_hosp.diagnoses_icd" in sql
assert "icd_code LIKE 'I10%'" in sql
def test_multiple_icd_codes(self):
"""Multiple ICD codes should create OR conditions."""
criteria = QueryCohortInput(icd_codes=["I10", "E11"])
sql = build_cohort_count_sql(criteria)
assert "icd_code LIKE 'I10%'" in sql
assert "icd_code LIKE 'E11%'" in sql
assert " OR " in sql
def test_icd_code_with_dot(self):
"""ICD code with dot should be valid."""
criteria = QueryCohortInput(icd_codes=["E11.9"])
sql = build_cohort_count_sql(criteria)
assert "icd_code LIKE 'E11.9%'" in sql
def test_icd_match_all_single_code(self):
"""Single ICD code with match_all should create single EXISTS."""
criteria = QueryCohortInput(icd_codes=["I10"], icd_match_all=True)
sql = build_cohort_count_sql(criteria)
assert "EXISTS" in sql
assert "icd_code LIKE 'I10%'" in sql
assert " OR " not in sql
def test_icd_match_all_multiple_codes(self):
"""Multiple ICD codes with match_all should create separate EXISTS for each."""
criteria = QueryCohortInput(icd_codes=["I10", "E11"], icd_match_all=True)
sql = build_cohort_count_sql(criteria)
# Each code should have its own EXISTS clause
assert sql.count("EXISTS") == 2
assert "icd_code LIKE 'I10%'" in sql
assert "icd_code LIKE 'E11%'" in sql
# Should not use OR to combine them
assert " OR " not in sql or "OR" not in sql.split("diagnoses_icd")[1]
def test_icd_match_any_is_default(self):
"""Without icd_match_all, should use OR (default behavior)."""
criteria = QueryCohortInput(icd_codes=["I10", "E11"])
sql = build_cohort_count_sql(criteria)
# Should have single EXISTS with OR
assert sql.count("EXISTS") == 1
assert " OR " in sql
class TestIcdCodesValidation:
"""Test ICD codes validation errors."""
def test_icd_injection_semicolon(self):
"""ICD code with semicolon should raise ValueError."""
criteria = QueryCohortInput(icd_codes=["I10; DROP TABLE"])
with pytest.raises(ValueError, match="Invalid ICD code format"):
build_cohort_count_sql(criteria)
def test_icd_injection_quote(self):
"""ICD code with quote should raise ValueError."""
criteria = QueryCohortInput(icd_codes=["I10'"])
with pytest.raises(ValueError, match="Invalid ICD code format"):
build_cohort_count_sql(criteria)
def test_icd_injection_dash(self):
"""ICD code with SQL comment should raise ValueError."""
criteria = QueryCohortInput(icd_codes=["I10--"])
with pytest.raises(ValueError, match="Invalid ICD code format"):
build_cohort_count_sql(criteria)
def test_icd_empty_string(self):
"""Empty ICD code string should raise ValueError."""
criteria = QueryCohortInput(icd_codes=[""])
with pytest.raises(ValueError, match="cannot be empty"):
build_cohort_count_sql(criteria)
def test_icd_empty_list_is_valid(self):
"""Empty ICD codes list should not add condition."""
criteria = QueryCohortInput(icd_codes=[])
sql = build_cohort_count_sql(criteria)
assert "diagnoses_icd" not in sql
class TestIcuStayCriteria:
"""Test ICU stay filtering criteria."""
def test_has_icu_stay_true(self):
"""has_icu_stay=True should JOIN icustays and include icu_stay_count."""
criteria = QueryCohortInput(has_icu_stay=True)
sql = build_cohort_count_sql(criteria)
# Should JOIN icustays table (not EXISTS) for ICU stay count
assert "JOIN mimiciv_icu.icustays" in sql
assert "icu_stay_count" in sql
assert "COUNT(DISTINCT i.stay_id)" in sql
assert "NOT EXISTS" not in sql
def test_has_icu_stay_false(self):
"""has_icu_stay=False should add NOT EXISTS subquery."""
criteria = QueryCohortInput(has_icu_stay=False)
sql = build_cohort_count_sql(criteria)
assert "NOT EXISTS" in sql
assert "mimiciv_icu.icustays" in sql
class TestMortalityCriteria:
"""Test in-hospital mortality filtering criteria."""
def test_mortality_true(self):
"""in_hospital_mortality=True should filter for deaths."""
criteria = QueryCohortInput(in_hospital_mortality=True)
sql = build_cohort_count_sql(criteria)
assert "hospital_expire_flag = 1" in sql
def test_mortality_false(self):
"""in_hospital_mortality=False should filter for survivors."""
criteria = QueryCohortInput(in_hospital_mortality=False)
sql = build_cohort_count_sql(criteria)
assert "hospital_expire_flag = 0" in sql
class TestCombinedCriteria:
"""Test multiple criteria combined."""
def test_age_and_gender(self):
"""Age and gender criteria combined."""
criteria = QueryCohortInput(age_min=18, age_max=65, gender="M")
sql = build_cohort_count_sql(criteria)
assert "p.anchor_age >= 18" in sql
assert "p.anchor_age <= 65" in sql
assert "p.gender = 'M'" in sql
def test_all_criteria(self):
"""All criteria combined."""
criteria = QueryCohortInput(
age_min=18,
age_max=80,
gender="F",
icd_codes=["I10", "E11"],
has_icu_stay=True,
in_hospital_mortality=False,
)
sql = build_cohort_count_sql(criteria)
assert "p.anchor_age >= 18" in sql
assert "p.anchor_age <= 80" in sql
assert "p.gender = 'F'" in sql
assert "diagnoses_icd" in sql
assert "icd_code LIKE 'I10%'" in sql
assert "mimiciv_icu.icustays" in sql
assert "EXISTS" in sql
assert "hospital_expire_flag = 0" in sql
def test_combined_criteria_pass_validation(self):
"""Combined criteria SQL should pass safety validation."""
criteria = QueryCohortInput(
age_min=18,
age_max=80,
gender="F",
icd_codes=["I10", "E11.9", "J18"],
has_icu_stay=True,
in_hospital_mortality=False,
)
for builder in [
build_cohort_count_sql,
build_cohort_demographics_sql,
build_gender_distribution_sql,
]:
sql = builder(criteria)
safe, msg = is_safe_query(sql)
assert safe, f"SQL failed validation: {msg}"
class TestSqlSafety:
"""Test that all generated SQL passes is_safe_query()."""
@pytest.mark.parametrize(
"criteria",
[
QueryCohortInput(),
QueryCohortInput(age_min=0),
QueryCohortInput(age_max=130),
QueryCohortInput(gender="M"),
QueryCohortInput(icd_codes=["A00"]),
QueryCohortInput(icd_codes=["A00.1", "B99.9", "Z87"]),
QueryCohortInput(icd_codes=["I10", "E11"], icd_match_all=True),
QueryCohortInput(icd_codes=["I10", "E11"], icd_match_all=False),
QueryCohortInput(has_icu_stay=True),
QueryCohortInput(has_icu_stay=False),
QueryCohortInput(in_hospital_mortality=True),
QueryCohortInput(in_hospital_mortality=False),
QueryCohortInput(
age_min=18,
age_max=65,
gender="F",
icd_codes=["I10"],
has_icu_stay=True,
in_hospital_mortality=True,
),
QueryCohortInput(
age_min=18,
age_max=65,
gender="F",
icd_codes=["I10", "E11"],
icd_match_all=True,
has_icu_stay=True,
in_hospital_mortality=True,
),
],
)
def test_sql_passes_safety_check(self, criteria):
"""All valid criteria should produce safe SQL."""
for builder in [
build_cohort_count_sql,
build_cohort_demographics_sql,
build_gender_distribution_sql,
]:
sql = builder(criteria)
safe, msg = is_safe_query(sql)
assert safe, f"SQL failed validation for {criteria}: {msg}\nSQL: {sql}"
