111 lines
5.2 KiB
Python
111 lines
5.2 KiB
Python
import sys
|
|
import os
|
|
import logging
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Add backend path & activate venv if possible
|
|
sys.path.insert(0, "/app")
|
|
|
|
from backend.services.classification import ClassificationService
|
|
from backend.database import Company, Industry
|
|
|
|
# Setup basic logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
class TestOpenerGeneration(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
"""Set up a mock environment."""
|
|
self.service = ClassificationService()
|
|
|
|
# Mock a database session
|
|
self.mock_db = MagicMock()
|
|
|
|
# Mock Industry object (as if read from DB)
|
|
self.mock_industry = Industry(
|
|
name="Leisure - Wet & Spa",
|
|
pains="Hohes Unfallrisiko durch Nässe, strenge Hygiene-Anforderungen.",
|
|
ops_focus_secondary=False
|
|
)
|
|
|
|
# Mock Company object
|
|
self.mock_company = Company(
|
|
id=1,
|
|
name="Therme Erding",
|
|
industry_ai="Leisure - Wet & Spa"
|
|
)
|
|
# Add the fields we are testing
|
|
self.mock_company.ai_opener = None
|
|
self.mock_company.ai_opener_secondary = None
|
|
|
|
@patch('backend.services.classification.call_gemini_flash')
|
|
@patch('backend.services.classification.ClassificationService._run_llm_classification_prompt')
|
|
def test_dual_opener_generation(self, mock_classification_call, mock_gemini_call):
|
|
"""
|
|
Test that both primary and secondary openers are generated and stored.
|
|
"""
|
|
print("\n--- Running Integration Test for Dual Opener Generation ---")
|
|
|
|
# --- Configure Mocks ---
|
|
# 1. Mock the classification call to return the correct industry
|
|
mock_classification_call.return_value = "Leisure - Wet & Spa"
|
|
|
|
# 2. Mock the opener generation calls (Gemini)
|
|
mock_gemini_call.side_effect = [
|
|
"Der reibungslose Betrieb ist entscheidend, um maximale Sicherheit zu gewährleisten.", # Mocked Primary
|
|
"Ein einzigartiges Gästeerlebnis ist der Schlüssel zum Erfolg." # Mocked Secondary
|
|
]
|
|
|
|
# Mock the content loader to return some text
|
|
with patch.object(self.service, '_get_website_content_and_url', return_value=("Die Therme Erding ist die größte Therme der Welt.", "http://mock.com")):
|
|
# --- Execute the Method ---
|
|
print("1. Calling classify_company_potential...")
|
|
# We patch the metric extraction to isolate the opener logic
|
|
with patch.object(self.service, 'extract_metrics_for_industry', return_value=self.mock_company):
|
|
# The method under test!
|
|
result_company = self.service.classify_company_potential(self.mock_company, self.mock_db)
|
|
|
|
# --- Assertions ---
|
|
print("2. Verifying results...")
|
|
|
|
# 1. Check that Gemini was called twice for the OPENERS
|
|
self.assertEqual(mock_gemini_call.call_count, 2, "❌ FAILED: AI model for OPENERS should have been called twice.")
|
|
print(" ✅ AI model for openers was called twice.")
|
|
|
|
# 2. Check that the classification prompt was called
|
|
self.assertEqual(mock_classification_call.call_count, 1, "❌ FAILED: Classification prompt should have been called once.")
|
|
print(" ✅ Classification prompt was called once.")
|
|
|
|
# 3. Check prompts contained the correct focus
|
|
first_call_args, _ = mock_gemini_call.call_args_list[0]
|
|
second_call_args, _ = mock_gemini_call.call_args_list[1]
|
|
self.assertIn("FOKUS: PRIMÄR-PROZESSE", first_call_args[0], "❌ FAILED: First call prompt did not have PRIMARY focus.")
|
|
print(" ✅ First opener call had PRIMARY focus.")
|
|
self.assertIn("FOKUS: SEKUNDÄR-PROZESSE", second_call_args[0], "❌ FAILED: Second call prompt did not have SECONDARY focus.")
|
|
print(" ✅ Second opener call had SECONDARY focus.")
|
|
|
|
# 4. Check that the results were stored on the company object
|
|
self.assertIsNotNone(result_company.ai_opener, "❌ FAILED: ai_opener (primary) was not set.")
|
|
self.assertIsNotNone(result_company.ai_opener_secondary, "❌ FAILED: ai_opener_secondary was not set.")
|
|
print(" ✅ Both ai_opener fields were set on the company object.")
|
|
|
|
# 5. Check content of the fields
|
|
self.assertIn("Sicherheit", result_company.ai_opener, "❌ FAILED: Primary opener content mismatch.")
|
|
print(f" -> Primary Opener: '{result_company.ai_opener}'")
|
|
self.assertIn("Gästeerlebnis", result_company.ai_opener_secondary, "❌ FAILED: Secondary opener content mismatch.")
|
|
print(f" -> Secondary Opener: '{result_company.ai_opener_secondary}'")
|
|
|
|
print("\n--- ✅ PASSED: Dual Opener logic is working correctly. ---")
|
|
|
|
if __name__ == '__main__':
|
|
# Patch the _load_industry_definitions to return our mock
|
|
with patch('backend.services.classification.ClassificationService._load_industry_definitions') as mock_load:
|
|
# Provide a description, so the classifier can match the industry
|
|
mock_load.return_value = [Industry(
|
|
name="Leisure - Wet & Spa",
|
|
description="Thermalbad, Spa, Wasserwelten, Saunen und Rutschenparks.",
|
|
pains="Hygiene"
|
|
)]
|
|
unittest.main()
|