import unittest
from llm.chains import load_prompt_template, validate_response
from llm.models import AnalysisFlags

class PromptTests(unittest.TestCase):
    def test_prompt_loading(self):
        """Test that prompt template loads correctly"""
        try:
            template = load_prompt_template()
            self.assertIsNotNone(template)
            self.assertIn("issueKey", template.input_variables)
        except Exception as e:
            self.fail(f"Prompt loading failed: {str(e)}")

    def test_response_validation(self):
        """Test response validation logic"""
        valid_response = {
            "hasMultipleEscalations": False,
            "customerSentiment": "neutral"
        }
        invalid_response = {
            "customerSentiment": "neutral"
        }
        
        self.assertTrue(validate_response(valid_response))
        self.assertFalse(validate_response(invalid_response))

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