218 lines
6.7 KiB
Python
218 lines
6.7 KiB
Python
"""
|
|
Unit tests for content generation service
|
|
"""
|
|
|
|
import pytest
|
|
import json
|
|
from unittest.mock import Mock, MagicMock, patch
|
|
from src.generation.service import ContentGenerationService, GenerationError
|
|
from src.database.models import Project, GeneratedContent
|
|
from src.generation.rule_engine import ValidationResult
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_session():
|
|
return Mock()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_config():
|
|
config = Mock()
|
|
config.ai_service.max_tokens = 4000
|
|
config.content_rules.cora_validation.round_averages_down = True
|
|
config.content_rules.cora_validation.tier_1_strict = True
|
|
return config
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_project():
|
|
project = Mock(spec=Project)
|
|
project.id = 1
|
|
project.main_keyword = "test keyword"
|
|
project.word_count = 1000
|
|
project.term_frequency = 3
|
|
project.tier = 1
|
|
project.h2_total = 5
|
|
project.h2_exact = 1
|
|
project.h2_related_search = 1
|
|
project.h2_entities = 2
|
|
project.h3_total = 10
|
|
project.h3_exact = 1
|
|
project.h3_related_search = 2
|
|
project.h3_entities = 3
|
|
project.entities = ["entity1", "entity2", "entity3"]
|
|
project.related_searches = ["search1", "search2", "search3"]
|
|
return project
|
|
|
|
|
|
@pytest.fixture
|
|
def service(mock_session, mock_config):
|
|
with patch('src.generation.service.AIClient'):
|
|
service = ContentGenerationService(mock_session, mock_config)
|
|
return service
|
|
|
|
|
|
def test_service_initialization(service):
|
|
"""Test service initializes correctly"""
|
|
assert service.session is not None
|
|
assert service.config is not None
|
|
assert service.ai_client is not None
|
|
assert service.content_repo is not None
|
|
assert service.rule_engine is not None
|
|
|
|
|
|
def test_generate_title_success(service, mock_project):
|
|
"""Test successful title generation"""
|
|
service.ai_client.generate = Mock(return_value="Test Keyword Complete Guide")
|
|
service.validator.validate_title = Mock(return_value=(True, []))
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.title_attempts = 0
|
|
service.content_repo.update = Mock()
|
|
|
|
result = service._generate_title(mock_project, content_record, "test-model", 3)
|
|
|
|
assert result == "Test Keyword Complete Guide"
|
|
assert service.ai_client.generate.called
|
|
|
|
|
|
def test_generate_title_validation_retry(service, mock_project):
|
|
"""Test title generation retries on validation failure"""
|
|
service.ai_client.generate = Mock(side_effect=[
|
|
"Wrong Title",
|
|
"Test Keyword Guide"
|
|
])
|
|
service.validator.validate_title = Mock(side_effect=[
|
|
(False, ["Missing keyword"]),
|
|
(True, [])
|
|
])
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.title_attempts = 0
|
|
service.content_repo.update = Mock()
|
|
|
|
result = service._generate_title(mock_project, content_record, "test-model", 3)
|
|
|
|
assert result == "Test Keyword Guide"
|
|
assert service.ai_client.generate.call_count == 2
|
|
|
|
|
|
def test_generate_title_max_retries_exceeded(service, mock_project):
|
|
"""Test title generation fails after max retries"""
|
|
service.ai_client.generate = Mock(return_value="Wrong Title")
|
|
service.validator.validate_title = Mock(return_value=(False, ["Missing keyword"]))
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.title_attempts = 0
|
|
service.content_repo.update = Mock()
|
|
|
|
with pytest.raises(GenerationError, match="validation failed"):
|
|
service._generate_title(mock_project, content_record, "test-model", 2)
|
|
|
|
|
|
def test_generate_outline_success(service, mock_project):
|
|
"""Test successful outline generation"""
|
|
outline_data = {
|
|
"h1": "Test Keyword Overview",
|
|
"sections": [
|
|
{"h2": "Test Keyword Basics", "h3s": ["Sub 1", "Sub 2"]},
|
|
{"h2": "Advanced Topics", "h3s": ["Sub 3"]}
|
|
]
|
|
}
|
|
|
|
service.ai_client.generate_json = Mock(return_value=outline_data)
|
|
service.validator.validate_outline = Mock(return_value=(True, [], {}))
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.outline_attempts = 0
|
|
service.content_repo.update = Mock()
|
|
|
|
result = service._generate_outline(
|
|
mock_project, "Test Title", content_record, "test-model", 3
|
|
)
|
|
|
|
assert result == outline_data
|
|
assert service.ai_client.generate_json.called
|
|
|
|
|
|
def test_generate_outline_with_augmentation(service, mock_project):
|
|
"""Test outline generation with programmatic augmentation"""
|
|
initial_outline = {
|
|
"h1": "Test Keyword Overview",
|
|
"sections": [
|
|
{"h2": "Introduction", "h3s": []}
|
|
]
|
|
}
|
|
|
|
augmented_outline = {
|
|
"h1": "Test Keyword Overview",
|
|
"sections": [
|
|
{"h2": "Test Keyword Introduction", "h3s": ["Sub 1"]},
|
|
{"h2": "Advanced Topics", "h3s": []}
|
|
]
|
|
}
|
|
|
|
service.ai_client.generate_json = Mock(return_value=initial_outline)
|
|
service.validator.validate_outline = Mock(side_effect=[
|
|
(False, ["Not enough H2s"], {"h2_exact": 1}),
|
|
(True, [], {})
|
|
])
|
|
service.augmenter.augment_outline = Mock(return_value=(augmented_outline, {}))
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.outline_attempts = 0
|
|
content_record.augmented = False
|
|
service.content_repo.update = Mock()
|
|
|
|
result = service._generate_outline(
|
|
mock_project, "Test Title", content_record, "test-model", 3
|
|
)
|
|
|
|
assert service.augmenter.augment_outline.called
|
|
|
|
|
|
def test_generate_content_success(service, mock_project):
|
|
"""Test successful content generation"""
|
|
html_content = "<h1>Test</h1><p>Content</p>"
|
|
|
|
service.ai_client.generate = Mock(return_value=html_content)
|
|
|
|
validation_result = Mock(spec=ValidationResult)
|
|
validation_result.passed = True
|
|
validation_result.errors = []
|
|
validation_result.warnings = []
|
|
validation_result.to_dict = Mock(return_value={})
|
|
|
|
service.validator.validate_content = Mock(return_value=(True, validation_result))
|
|
|
|
content_record = Mock(spec=GeneratedContent)
|
|
content_record.content_attempts = 0
|
|
service.content_repo.update = Mock()
|
|
|
|
outline = {"h1": "Test", "sections": []}
|
|
|
|
result = service._generate_content(
|
|
mock_project, "Test Title", outline, content_record, "test-model", 3
|
|
)
|
|
|
|
assert result == html_content
|
|
|
|
|
|
def test_format_outline_for_prompt(service):
|
|
"""Test outline formatting for content prompt"""
|
|
outline = {
|
|
"h1": "Main Heading",
|
|
"sections": [
|
|
{"h2": "Section 1", "h3s": ["Sub 1", "Sub 2"]},
|
|
{"h2": "Section 2", "h3s": ["Sub 3"]}
|
|
]
|
|
}
|
|
|
|
result = service._format_outline_for_prompt(outline)
|
|
|
|
assert "H1: Main Heading" in result
|
|
assert "H2: Section 1" in result
|
|
assert "H3: Sub 1" in result
|
|
assert "H2: Section 2" in result
|
|
|