|
|
from unittest.mock import AsyncMock, MagicMock |
|
|
|
|
|
import pytest |
|
|
|
|
|
from vsp.app.classifiers.education_classifier import EducationClassification, EducationClassifier, SchoolType |
|
|
from vsp.app.model.linkedin.linkedin_models import DateComponent, Education, LinkedinProfile |
|
|
|
|
|
|
|
|
@pytest.fixture |
|
|
def mock_llm_service(): |
|
|
return AsyncMock() |
|
|
|
|
|
|
|
|
@pytest.fixture |
|
|
def mock_prompt_loader(): |
|
|
loader = MagicMock() |
|
|
loader.load_template.return_value = MagicMock() |
|
|
loader.create_prompt.return_value = AsyncMock() |
|
|
return loader |
|
|
|
|
|
|
|
|
@pytest.fixture |
|
|
def education_classifier(mock_llm_service, mock_prompt_loader): |
|
|
return EducationClassifier(llm_service=mock_llm_service, prompt_loader=mock_prompt_loader) |
|
|
|
|
|
|
|
|
@pytest.fixture |
|
|
def sample_linkedin_profile(): |
|
|
return LinkedinProfile( |
|
|
first_name="John", |
|
|
last_name="Doe", |
|
|
educations=[ |
|
|
Education( |
|
|
school_name="Stanford University", |
|
|
degree="Master of Business Administration", |
|
|
field_of_study="Business Administration", |
|
|
start=DateComponent(year=2018), |
|
|
end=DateComponent(year=2020), |
|
|
) |
|
|
], |
|
|
) |
|
|
|
|
|
|
|
|
@pytest.mark.asyncio |
|
|
async def test_classify_education(education_classifier, sample_linkedin_profile, mock_prompt_loader): |
|
|
mock_prompt = mock_prompt_loader.create_prompt.return_value |
|
|
mock_prompt.evaluate.return_value = EducationClassification( |
|
|
output=SchoolType.MBA, |
|
|
confidence=0.95, |
|
|
reasoning="This is a Master of Business Administration degree from Stanford University.", |
|
|
) |
|
|
|
|
|
result = await education_classifier.classify_education( |
|
|
sample_linkedin_profile, sample_linkedin_profile.educations[0] |
|
|
) |
|
|
|
|
|
assert isinstance(result, EducationClassification) |
|
|
assert result.output == SchoolType.MBA |
|
|
assert result.confidence == 0.95 |
|
|
assert "Master of Business Administration" in result.reasoning |
|
|
|
|
|
|
|
|
@pytest.mark.parametrize( |
|
|
"output,expected", |
|
|
[ |
|
|
("PRIMARY_SECONDARY", SchoolType.PRIMARY_SECONDARY), |
|
|
("UNDERGRAD_INCOMPLETE", SchoolType.UNDERGRAD_INCOMPLETE), |
|
|
("UNDERGRAD_COMPLETED", SchoolType.UNDERGRAD_COMPLETED), |
|
|
("MBA", SchoolType.MBA), |
|
|
("LAW_SCHOOL", SchoolType.LAW_SCHOOL), |
|
|
("GRAD_SCHOOL", SchoolType.GRAD_SCHOOL), |
|
|
("PHD", SchoolType.PHD), |
|
|
("OTHER", SchoolType.OTHER), |
|
|
], |
|
|
) |
|
|
def test_parse_output(output, expected): |
|
|
parsed = EducationClassifier._parse_output(f"output: {output}\nconfidence: 0.9\nreasoning: Test reasoning") |
|
|
assert parsed.output == expected |
|
|
assert parsed.confidence == 0.9 |
|
|
assert parsed.reasoning == "Test reasoning" |
|
|
|
|
|
|
|
|
def test_parse_output_invalid(): |
|
|
with pytest.raises(ValueError): |
|
|
EducationClassifier._parse_output("output: INVALID\nconfidence: 0.9\nreasoning: Test reasoning") |
|
|
|
