File size: 2,829 Bytes
3b993c4 145d03b 24d33b9 3b993c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
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")
|