Spaces:
Runtime error
Runtime error
| """Tests for risk aggregation utilities.""" | |
| import pytest | |
| from sentinel.models import RiskScore | |
| from sentinel.risk_aggregation import ( | |
| format_scores_for_llm, | |
| format_scores_for_pdf, | |
| group_scores_by_cancer_type, | |
| ) | |
| def test_group_scores_by_cancer_type(): | |
| """Test grouping risk scores by cancer type.""" | |
| scores = [ | |
| RiskScore( | |
| name="Gail Model", | |
| score="5%", | |
| cancer_type="Breast Cancer", | |
| description="5-year risk", | |
| ), | |
| RiskScore( | |
| name="Claus Model", | |
| score="3%", | |
| cancer_type="Breast Cancer", | |
| description="Lifetime risk", | |
| ), | |
| RiskScore( | |
| name="PLCOm2012", | |
| score="2%", | |
| cancer_type="Lung Cancer", | |
| description="6-year risk", | |
| ), | |
| ] | |
| grouped = group_scores_by_cancer_type(scores) | |
| assert len(grouped) == 2 | |
| assert "Breast Cancer" in grouped | |
| assert "Lung Cancer" in grouped | |
| assert len(grouped["Breast Cancer"]) == 2 | |
| assert len(grouped["Lung Cancer"]) == 1 | |
| assert grouped["Breast Cancer"][0].name == "Gail Model" | |
| assert grouped["Breast Cancer"][1].name == "Claus Model" | |
| def test_group_scores_empty(): | |
| """Test grouping with empty list.""" | |
| grouped = group_scores_by_cancer_type([]) | |
| assert grouped == {} | |
| def test_group_scores_no_cancer_type(): | |
| """Test grouping with scores that have no cancer type.""" | |
| scores = [ | |
| RiskScore(name="Test Model", score="5%", cancer_type=None), | |
| ] | |
| with pytest.raises(ValueError, match=r"Test Model.*missing cancer_type"): | |
| group_scores_by_cancer_type(scores) | |
| def test_format_scores_for_llm(): | |
| """Test formatting scores for LLM context.""" | |
| scores = [ | |
| RiskScore( | |
| name="Gail Model", | |
| score="5%", | |
| cancer_type="Breast Cancer", | |
| description="5-year risk", | |
| interpretation="Low to moderate risk", | |
| references=["Gail et al., 1989"], | |
| ), | |
| ] | |
| grouped = group_scores_by_cancer_type(scores) | |
| formatted = format_scores_for_llm(grouped) | |
| assert "# Calculated Risk Scores (Ground Truth)" in formatted | |
| assert "Breast Cancer" in formatted | |
| assert "Gail Model" in formatted | |
| assert "5%" in formatted | |
| assert "5-year risk" in formatted | |
| assert "Low to moderate risk" in formatted | |
| assert "Gail et al., 1989" in formatted | |
| assert "DO NOT generate your own risk levels" in formatted | |
| def test_format_scores_for_llm_empty(): | |
| """Test formatting empty scores for LLM.""" | |
| formatted = format_scores_for_llm({}) | |
| assert formatted == "No risk scores calculated." | |
| def test_format_scores_for_pdf(): | |
| """Test formatting scores for PDF presentation.""" | |
| scores = [ | |
| RiskScore( | |
| name="Gail Model", | |
| score="5%", | |
| cancer_type="Breast Cancer", | |
| ), | |
| RiskScore( | |
| name="PLCOm2012", | |
| score="2%", | |
| cancer_type="Lung Cancer", | |
| ), | |
| ] | |
| grouped = group_scores_by_cancer_type(scores) | |
| formatted = format_scores_for_pdf(grouped) | |
| assert len(formatted) == 2 | |
| assert formatted[0][0] == "Breast Cancer" | |
| assert formatted[1][0] == "Lung Cancer" | |
| assert len(formatted[0][1]) == 1 | |
| assert len(formatted[1][1]) == 1 | |