Spaces:
Sleeping
Sleeping
| """Tests for improved hierarchical classification.""" | |
| import pytest | |
| import os | |
| from core.classification import ImprovedHierarchicalClassifier | |
| def test_classifier_initialization(): | |
| """Test classifier initialization.""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=False) | |
| assert classifier.hierarchy_name == "hospital" | |
| assert classifier.hierarchy is not None | |
| assert len(classifier.level1_keywords) > 0 | |
| def test_keyword_classification(): | |
| """Test keyword-based classification.""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=False) | |
| text = "Patient admission procedures require proper documentation and identification verification." | |
| result = classifier.classify_text(text) | |
| assert "level1" in result | |
| assert "level2" in result | |
| assert "level3" in result | |
| assert "doc_type" in result | |
| assert result["method"] == "keyword" | |
| def test_llm_classification(): | |
| """Test LLM-based classification (requires API key).""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=True) | |
| text = """ | |
| Patient Admission Protocol | |
| All patients must present valid identification upon admission. | |
| Emergency cases follow expedited procedures. | |
| """ | |
| result = classifier.classify_text(text) | |
| assert "level1" in result | |
| assert "level2" in result | |
| assert "level3" in result | |
| assert "doc_type" in result | |
| assert "confidence" in result | |
| assert 0 <= result["confidence"] <= 1 | |
| def test_doc_type_inference(): | |
| """Test document type inference.""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=False) | |
| policy_text = "This policy outlines the requirements for patient care." | |
| manual_text = "This manual provides step-by-step instructions for procedures." | |
| report_text = "This report summarizes the findings of our quality analysis." | |
| policy_result = classifier._infer_doc_type(policy_text.lower()) | |
| manual_result = classifier._infer_doc_type(manual_text.lower()) | |
| report_result = classifier._infer_doc_type(report_text.lower()) | |
| assert policy_result == "policy" | |
| assert manual_result == "manual" | |
| assert report_result == "report" | |
| def test_classification_fallback(): | |
| """Test fallback when LLM fails.""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=False) | |
| # Even with empty text, should return valid structure | |
| result = classifier._fallback_classification("") | |
| assert "level1" in result | |
| assert "level2" in result | |
| assert "level3" in result | |
| assert "doc_type" in result | |
| assert result["confidence"] == 0.3 | |
| assert result["method"] == "keyword" | |
| def test_classification_with_override(): | |
| """Test classification with doc_type override.""" | |
| classifier = ImprovedHierarchicalClassifier("hospital", use_llm=False) | |
| text = "Some document about patient care." | |
| result = classifier.classify_text(text, doc_type="manual") | |
| assert result["doc_type"] == "manual" |