Add a test for main.py (#3)

* Add a test for main.py

* Improved nesting on the test classes

src/vsp/app/classifiers/education_classifier.py

@@ -51,6 +51,7 @@ class EducationClassification(BaseModel):
     output: SchoolType = Field(description="The classified school type")
     confidence: float = Field(description="Confidence level between 0.0 and 1.0")
     reasoning: str = Field(description="Explanation for the classification")
+    model_config = {"frozen": True}  # This makes the model immutable and hashable
 
 
 class EducationClassifier:

src/vsp/app/classifiers/work_experience/general_work_experience_classifier.py

@@ -73,6 +73,7 @@ class WorkExperienceClassification(BaseModel):
     secondary_job_type: SecondaryJobType = Field(description="The classified secondary job type")
     confidence: float = Field(description="Confidence level between 0.0 and 1.0", ge=0.0, le=1.0)
     reasoning: str = Field(description="Explanation for the classification")
+    model_config = {"frozen": True}  # This makes the model immutable and hashable
 
 
 class WorkExperienceClassifier:

src/vsp/app/classifiers/work_experience/investment_banking_group_classifier.py

@@ -70,6 +70,7 @@ class InvestmentBankingGroupClassification(BaseModel):
     investment_banking_group: InvestmentBankingGroup = Field(description="The investment banking group")
     confidence: float = Field(description="Confidence level between 0.0 and 1.0", ge=0.0, le=1.0)
     reasoning: str = Field(description="Explanation for the classification")
+    model_config = {"frozen": True}  # This makes the model immutable and hashable
 
 
 class InvestmentBankingGroupClassifier:

tests/vsp/app/test_main.py

@@ -0,0 +1,122 @@
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from vsp.app.classifiers.education_classifier import EducationClassification, SchoolType
+from vsp.app.classifiers.work_experience.general_work_experience_classifier import (
+    PrimaryJobType,
+    SecondaryJobType,
+    WorkExperienceClassification,
+)
+from vsp.app.classifiers.work_experience.investment_banking_group_classifier import (
+    InvestmentBankingGroup,
+    InvestmentBankingGroupClassification,
+)
+from vsp.app.main import LinkedinProfileClassificationResults, process_linkedin_profile
+from vsp.app.model.linkedin.linkedin_models import Education, LinkedinProfile, Position
+
+
+@pytest.fixture
+def sample_linkedin_profile():
+    return LinkedinProfile(
+        first_name="John",
+        last_name="Doe",
+        educations=[
+            Education(
+                school_name="Test University",
+                degree="MBA",
+                field_of_study="Business",
+            )
+        ],
+        position=[
+            Position(
+                title="Investment Banking Analyst",
+                company_name="Bank Corp",
+            ),
+            Position(
+                title="Software Engineer",
+                company_name="Tech Corp",
+            ),
+        ],
+    )
+
+
+@pytest.mark.asyncio
+async def test_process_linkedin_profile_investment_banking(sample_linkedin_profile):
+    with (
+        patch("vsp.app.main.EducationClassifier") as mock_education_classifier,
+        patch("vsp.app.main.WorkExperienceClassifier") as mock_work_experience_classifier,
+        patch("vsp.app.main.InvestmentBankingGroupClassifier") as mock_investment_banking_group_classifier,
+    ):
+        mock_education_classifier.return_value.classify_education = AsyncMock(
+            return_value=EducationClassification(output=SchoolType.MBA, confidence=1.0, reasoning="Test")
+        )
+        mock_work_experience_classifier.return_value.classify_work_experience = AsyncMock(
+            return_value=WorkExperienceClassification(
+                primary_job_type=PrimaryJobType.FULL_TIME,
+                secondary_job_type=SecondaryJobType.INVESTMENT_BANKING,
+                confidence=1.0,
+                reasoning="Test",
+            )
+        )
+        mock_investment_banking_group_classifier.return_value.classify_investment_banking_group = AsyncMock(
+            return_value=InvestmentBankingGroupClassification(
+                investment_banking_group=InvestmentBankingGroup.M_AND_A, confidence=1.0, reasoning="Test"
+            )
+        )
+
+        result = await process_linkedin_profile(sample_linkedin_profile)
+
+        assert isinstance(result, LinkedinProfileClassificationResults)
+        assert len(result.classified_educations) == 1
+        assert len(result.classified_work_experiences) == 2
+
+        # Check that the investment banking position has an investment banking group classification
+        ib_experience = result.classified_work_experiences[0]
+        assert ib_experience.work_experience_classification.secondary_job_type == SecondaryJobType.INVESTMENT_BANKING
+        assert ib_experience.investment_banking_classification is not None
+        assert (
+            ib_experience.investment_banking_classification.investment_banking_group == InvestmentBankingGroup.M_AND_A
+        )
+
+        # Check that the investment banking group classifier was called
+        mock_investment_banking_group_classifier.return_value.classify_investment_banking_group.assert_called()
+
+
+@pytest.mark.asyncio
+async def test_process_linkedin_profile_non_investment_banking(sample_linkedin_profile):
+    with (
+        patch("vsp.app.main.EducationClassifier") as mock_education_classifier,
+        patch("vsp.app.main.WorkExperienceClassifier") as mock_work_experience_classifier,
+        patch("vsp.app.main.InvestmentBankingGroupClassifier") as mock_investment_banking_group_classifier,
+    ):
+        mock_education_classifier.return_value.classify_education = AsyncMock(
+            return_value=EducationClassification(output=SchoolType.MBA, confidence=1.0, reasoning="Test")
+        )
+        mock_work_experience_classifier.return_value.classify_work_experience = AsyncMock(
+            return_value=WorkExperienceClassification(
+                primary_job_type=PrimaryJobType.FULL_TIME,
+                secondary_job_type=SecondaryJobType.ENGINEERING,
+                confidence=1.0,
+                reasoning="Test",
+            )
+        )
+        mock_investment_banking_group_classifier.return_value.classify_investment_banking_group = AsyncMock()
+
+        result = await process_linkedin_profile(sample_linkedin_profile)
+
+        assert isinstance(result, LinkedinProfileClassificationResults)
+        assert len(result.classified_educations) == 1
+        assert len(result.classified_work_experiences) == 2
+
+        # Check that none of the work experiences have an investment banking group classification
+        for experience in result.classified_work_experiences:
+            assert experience.work_experience_classification.secondary_job_type == SecondaryJobType.ENGINEERING
+            assert experience.investment_banking_classification is None
+
+        # Check that the investment banking group classifier was not called
+        mock_investment_banking_group_classifier.return_value.classify_investment_banking_group.assert_not_called()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])