tests/test_role_classifier.py

import unittest

from src.agents.role_utils import classify_role


class TestRoleClassifier(unittest.TestCase):
    def test_caregiver_query_maps_to_caregiver(self):
        self.assertEqual(classify_role("I am caring for my mother with diabetes"), "caregiver")

    def test_dietary_query_maps_to_dietary(self):
        self.assertEqual(classify_role("Can you suggest a low-carb meal plan?"), "dietary")


if __name__ == "__main__":
    unittest.main()