Spaces:
Sleeping
Sleeping
| """ | |
| Unit tests for Medical Diagnostic Environment | |
| Run with: python -m pytest tests/test_environment.py -v | |
| """ | |
| import pytest | |
| from server.environment import MedicalDiagnosticEnvironment | |
| from server.medical_data import ( | |
| PATIENT_CASES, | |
| calculate_question_reward, | |
| calculate_test_reward, | |
| calculate_diagnosis_accuracy, | |
| ) | |
| from models import DiagnosticAction | |
| class TestMedicalDiagnosticEnvironment: | |
| """Test suite for MedicalDiagnosticEnvironment""" | |
| def env(self): | |
| """Create a fresh environment for each test""" | |
| return MedicalDiagnosticEnvironment() | |
| def test_environment_initialization(self, env): | |
| """Test that environment initializes correctly""" | |
| assert env is not None | |
| assert hasattr(env, "reset") | |
| assert hasattr(env, "step") | |
| assert hasattr(env, "state") | |
| def test_reset_easy(self, env): | |
| """Test reset with easy difficulty""" | |
| observation = env.reset(difficulty="easy") | |
| assert observation is not None | |
| assert hasattr(observation, "message") | |
| assert "presentation" in observation.message.lower() or "patient" in observation.message.lower() | |
| assert env.current_case_id is not None | |
| assert env.current_difficulty == "easy" | |
| def test_reset_medium(self, env): | |
| """Test reset with medium difficulty""" | |
| observation = env.reset(difficulty="medium") | |
| assert observation is not None | |
| assert env.current_difficulty == "medium" | |
| def test_reset_hard(self, env): | |
| """Test reset with hard difficulty""" | |
| observation = env.reset(difficulty="hard") | |
| assert observation is not None | |
| assert env.current_difficulty == "hard" | |
| def test_ask_question_action(self, env): | |
| """Test asking a question""" | |
| env.reset(difficulty="easy") | |
| action = DiagnosticAction( | |
| action_type="ask_question", | |
| question="Does the patient have a fever?" | |
| ) | |
| result = env.step(action) | |
| assert result is not None | |
| assert hasattr(result, "reward") | |
| assert result.reward >= 0 # Questions give non-negative reward | |
| def test_order_test_action(self, env): | |
| """Test ordering a test""" | |
| env.reset(difficulty="easy") | |
| action = DiagnosticAction( | |
| action_type="order_test", | |
| test_name="Complete Blood Count" | |
| ) | |
| result = env.step(action) | |
| assert result is not None | |
| assert hasattr(result, "reward") | |
| assert result.reward >= 0 # Tests give non-negative reward | |
| def test_submit_diagnosis_action(self, env): | |
| """Test submitting a diagnosis""" | |
| env.reset(difficulty="easy") | |
| action = DiagnosticAction( | |
| action_type="submit_diagnosis", | |
| diagnosis="Common Flu" | |
| ) | |
| result = env.step(action) | |
| assert result is not None | |
| assert hasattr(result, "reward") | |
| assert result.done is True # Episode should end on diagnosis | |
| def test_max_steps_enforcement(self, env): | |
| """Test that episodes end after max steps""" | |
| env.reset(difficulty="easy") | |
| for _ in range(15): # Max 15 steps | |
| action = DiagnosticAction( | |
| action_type="ask_question", | |
| question="Test question" | |
| ) | |
| result = env.step(action) | |
| if result.done: | |
| break | |
| assert result.done is True | |
| def test_episode_summary(self, env): | |
| """Test episode summary generation""" | |
| env.reset(difficulty="easy") | |
| action = DiagnosticAction( | |
| action_type="submit_diagnosis", | |
| diagnosis="Test Diagnosis" | |
| ) | |
| env.step(action) | |
| summary = env.get_episode_summary() | |
| assert summary is not None | |
| assert "case_id" in summary | |
| assert "difficulty" in summary | |
| assert "accuracy" in summary | |
| assert "total_reward" in summary | |
| assert "steps" in summary | |
| def test_state_property(self, env): | |
| """Test the state property""" | |
| env.reset(difficulty="easy") | |
| state = env.state | |
| assert state is not None | |
| assert hasattr(state, "patient_id") | |
| assert hasattr(state, "step_count") | |
| assert hasattr(state, "true_diagnosis") | |
| def test_concurrent_sessions(self): | |
| """Test that environment supports concurrent sessions""" | |
| env = MedicalDiagnosticEnvironment() | |
| assert env.SUPPORTS_CONCURRENT_SESSIONS is True | |
| def test_multiple_episodes(self, env): | |
| """Test running multiple episodes""" | |
| for difficulty in ["easy", "medium", "hard"]: | |
| observation = env.reset(difficulty=difficulty) | |
| assert observation is not None | |
| assert env.current_difficulty == difficulty | |
| class TestMedicalData: | |
| """Test suite for medical data functions""" | |
| def test_question_reward_calculation(self): | |
| """Test question reward calculation""" | |
| # This is case-specific, so we just verify the function works | |
| case_id = next(iter(PATIENT_CASES)) | |
| reward = calculate_question_reward( | |
| case_id=case_id, | |
| question="Does the patient have a fever?" | |
| ) | |
| assert 0.0 <= reward <= 1.0 | |
| def test_test_reward_calculation(self): | |
| """Test test reward calculation""" | |
| case_id = next(iter(PATIENT_CASES)) | |
| reward = calculate_test_reward( | |
| case_id=case_id, | |
| test_name="CBC" | |
| ) | |
| assert 0.0 <= reward <= 1.0 | |
| def test_diagnosis_accuracy_exact_match(self): | |
| """Test exact diagnosis match""" | |
| case_id = next(iter(PATIENT_CASES)) | |
| accuracy = calculate_diagnosis_accuracy( | |
| case_id=case_id, | |
| submitted_diagnosis=PATIENT_CASES[case_id].get("true_diagnosis", "") | |
| ) | |
| assert accuracy == 1.0 | |
| def test_diagnosis_accuracy_partial(self): | |
| """Test partial diagnosis accuracy""" | |
| case_id = next(iter(PATIENT_CASES)) | |
| accuracy = calculate_diagnosis_accuracy( | |
| case_id=case_id, | |
| submitted_diagnosis="Pneumonia" | |
| ) | |
| assert 0.0 <= accuracy <= 1.0 | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v", "-s"]) | |