94 lines
2.5 KiB
Python
94 lines
2.5 KiB
Python
"""
|
|
Unit tests for content augmenter
|
|
"""
|
|
|
|
import pytest
|
|
from src.generation.augmenter import ContentAugmenter
|
|
|
|
|
|
@pytest.fixture
|
|
def augmenter():
|
|
return ContentAugmenter()
|
|
|
|
|
|
def test_augment_outline_add_h2_keyword(augmenter):
|
|
"""Test adding keyword to H2 headings"""
|
|
outline = {
|
|
"h1": "Main Title",
|
|
"sections": [
|
|
{"h2": "Introduction", "h3s": []},
|
|
{"h2": "Advanced Topics", "h3s": []}
|
|
]
|
|
}
|
|
|
|
missing = {"h2_exact": 1}
|
|
|
|
result, log = augmenter.augment_outline(
|
|
outline, missing, "test keyword", [], []
|
|
)
|
|
|
|
assert "test keyword" in result["sections"][0]["h2"].lower()
|
|
assert log["headings_modified"] > 0
|
|
|
|
|
|
def test_augment_outline_add_h3_entities(augmenter):
|
|
"""Test adding entity-based H3 headings"""
|
|
outline = {
|
|
"h1": "Main Title",
|
|
"sections": [
|
|
{"h2": "Section 1", "h3s": []}
|
|
]
|
|
}
|
|
|
|
missing = {"h3_entities": 2}
|
|
entities = ["entity1", "entity2", "entity3"]
|
|
|
|
result, log = augmenter.augment_outline(
|
|
outline, missing, "keyword", entities, []
|
|
)
|
|
|
|
assert log["h3_added"] == 2
|
|
assert any("entity1" in h3.lower()
|
|
for s in result["sections"]
|
|
for h3 in s.get("h3s", []))
|
|
|
|
|
|
def test_augment_content_insert_keywords(augmenter):
|
|
"""Test inserting keywords into content"""
|
|
html = "<p>This is a paragraph with enough words to allow keyword insertion for testing purposes.</p>"
|
|
missing = {"keyword_mentions": 2}
|
|
|
|
result, log = augmenter.augment_content(
|
|
html, missing, "keyword", [], []
|
|
)
|
|
|
|
assert log["keywords_inserted"] > 0
|
|
assert "keyword" in result.lower()
|
|
|
|
|
|
def test_augment_content_insert_entities(augmenter):
|
|
"""Test inserting entities into content"""
|
|
html = "<p>This is a long paragraph with many words that allows us to insert various terms naturally.</p>"
|
|
missing = {"entity_mentions": 2}
|
|
entities = ["entity1", "entity2"]
|
|
|
|
result, log = augmenter.augment_content(
|
|
html, missing, "keyword", entities, []
|
|
)
|
|
|
|
assert log["entities_inserted"] > 0
|
|
|
|
|
|
def test_add_paragraph_with_terms(augmenter):
|
|
"""Test adding a new paragraph with specific terms"""
|
|
html = "<h1>Title</h1><p>Existing content</p>"
|
|
terms = ["term1", "term2", "term3"]
|
|
|
|
result = augmenter.add_paragraph_with_terms(
|
|
html, terms, "entity", "main keyword"
|
|
)
|
|
|
|
assert "term1" in result or "term2" in result or "term3" in result
|
|
assert "main keyword" in result
|
|
|