112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
import os
|
|
from collections import Counter
|
|
from datetime import date
|
|
from unittest.mock import patch, MagicMock
|
|
from io import BytesIO
|
|
from PIL import Image as PILImage
|
|
|
|
from src.cover import (
|
|
generate_programmatic_cover,
|
|
generate_ai_cover,
|
|
generate_cover,
|
|
_get_dominant_category,
|
|
CATEGORY_PROMPTS,
|
|
)
|
|
|
|
|
|
def test_programmatic_cover_is_portrait(tmp_path):
|
|
path = generate_programmatic_cover(
|
|
output_dir=str(tmp_path),
|
|
week_start=date(2026, 4, 6),
|
|
week_end=date(2026, 4, 12),
|
|
headlines=["Article One", "Article Two", "Article Three"],
|
|
categories=["Government", "Government", "Culture"],
|
|
)
|
|
|
|
assert os.path.exists(path)
|
|
img = PILImage.open(path)
|
|
assert img.format == "JPEG"
|
|
assert img.size == (480, 800)
|
|
|
|
|
|
def test_programmatic_cover_no_headlines(tmp_path):
|
|
path = generate_programmatic_cover(
|
|
output_dir=str(tmp_path),
|
|
week_start=date(2026, 4, 6),
|
|
week_end=date(2026, 4, 12),
|
|
headlines=[],
|
|
categories=[],
|
|
)
|
|
|
|
assert os.path.exists(path)
|
|
img = PILImage.open(path)
|
|
assert img.size == (480, 800)
|
|
|
|
|
|
def test_ai_cover_is_portrait(tmp_path):
|
|
fake_img = PILImage.new("RGB", (480, 800), color="blue")
|
|
buf = BytesIO()
|
|
fake_img.save(buf, format="JPEG")
|
|
fake_bytes = buf.getvalue()
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.content = fake_bytes
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
with patch("src.cover.requests.get", return_value=mock_response):
|
|
path = generate_ai_cover(
|
|
output_dir=str(tmp_path),
|
|
week_start=date(2026, 4, 6),
|
|
week_end=date(2026, 4, 12),
|
|
headlines=["Test Headline"],
|
|
categories=["Government"],
|
|
)
|
|
|
|
assert os.path.exists(path)
|
|
img = PILImage.open(path)
|
|
assert img.format == "JPEG"
|
|
assert img.size == (480, 800)
|
|
|
|
|
|
def test_ai_cover_falls_back_to_programmatic(tmp_path):
|
|
with patch("src.cover.requests.get", side_effect=Exception("API down")):
|
|
path = generate_ai_cover(
|
|
output_dir=str(tmp_path),
|
|
week_start=date(2026, 4, 6),
|
|
week_end=date(2026, 4, 12),
|
|
headlines=["Test"],
|
|
categories=["Government"],
|
|
)
|
|
|
|
assert os.path.exists(path)
|
|
img = PILImage.open(path)
|
|
assert img.format == "JPEG"
|
|
assert img.size == (480, 800)
|
|
|
|
|
|
def test_generate_cover_dispatches(tmp_path):
|
|
with patch("src.cover.generate_ai_cover") as mock_ai:
|
|
mock_ai.return_value = "/fake/ai.jpg"
|
|
result = generate_cover("ai", str(tmp_path), date(2026, 4, 6),
|
|
date(2026, 4, 12), ["A"], ["Government"])
|
|
assert result == "/fake/ai.jpg"
|
|
mock_ai.assert_called_once()
|
|
|
|
with patch("src.cover.generate_programmatic_cover") as mock_prog:
|
|
mock_prog.return_value = "/fake/text.jpg"
|
|
result = generate_cover("text", str(tmp_path), date(2026, 4, 6),
|
|
date(2026, 4, 12), ["A"], ["Government"])
|
|
assert result == "/fake/text.jpg"
|
|
mock_prog.assert_called_once()
|
|
|
|
|
|
def test_get_dominant_category():
|
|
assert _get_dominant_category(["Government", "Government", "Culture"]) == "Government"
|
|
assert _get_dominant_category(["Culture", "Government"]) == "Culture"
|
|
assert _get_dominant_category([]) == "Default"
|
|
|
|
|
|
def test_category_prompts_include_massachusetts():
|
|
for category, prompt in CATEGORY_PROMPTS.items():
|
|
assert "Massachusetts" in prompt, f"Prompt for {category} missing 'Massachusetts'"
|