kol-insight/backend/tests/test_export_api.py

import pytest
from io import BytesIO
from openpyxl import load_workbook

from app.services.export_service import generate_excel, generate_csv, COLUMN_HEADERS


class TestExportService:
    """Tests for Export Service."""

    @pytest.fixture
    def sample_export_data(self):
        """Sample data for export testing."""
        return [
            {
                "item_id": "item_001",
                "title": "测试视频1",
                "viral_type": "爆款",
                "video_url": "https://example.com/1",
                "star_id": "star_001",
                "star_unique_id": "unique_001",
                "star_nickname": "测试达人1",
                "publish_time": "2026-01-28T10:00:00",
                "natural_play_cnt": 100000,
                "heated_play_cnt": 50000,
                "total_play_cnt": 150000,
                "total_interact": 5000,
                "like_cnt": 3000,
                "share_cnt": 1000,
                "comment_cnt": 1000,
                "new_a3_rate": 0.05,
                "after_view_search_uv": 500,
                "return_search_cnt": 200,
                "industry_id": "ind_001",
                "industry_name": "美妆",
                "brand_id": "brand_001",
                "brand_name": "测试品牌",
                "estimated_video_cost": 10000.0,
                "estimated_natural_cpm": 100.0,
                "estimated_natural_search_uv": 333.33,
                "estimated_natural_search_cost": 30.0,
            }
        ]

    def test_generate_excel_success(self, sample_export_data):
        """Test Excel generation."""
        content = generate_excel(sample_export_data)

        assert content is not None
        assert len(content) > 0

        # 验证可以被 openpyxl 读取
        wb = load_workbook(BytesIO(content))
        ws = wb.active

        # 验证表头
        assert ws.cell(row=1, column=1).value == "视频ID"
        assert ws.cell(row=1, column=2).value == "视频标题"

        # 验证数据行
        assert ws.cell(row=2, column=1).value == "item_001"
        assert ws.cell(row=2, column=2).value == "测试视频1"

    def test_generate_excel_empty_data(self):
        """Test Excel generation with empty data."""
        content = generate_excel([])

        assert content is not None
        wb = load_workbook(BytesIO(content))
        ws = wb.active

        # 应该只有表头
        assert ws.max_row == 1

    def test_generate_csv_success(self, sample_export_data):
        """Test CSV generation."""
        content = generate_csv(sample_export_data)

        assert content is not None
        assert len(content) > 0

        # 验证 CSV 内容
        lines = content.decode("utf-8-sig").split("\n")
        assert len(lines) >= 2  # 表头 + 至少一行数据

        # 验证表头
        assert "视频ID" in lines[0]
        assert "视频标题" in lines[0]

    def test_generate_csv_empty_data(self):
        """Test CSV generation with empty data."""
        content = generate_csv([])

        assert content is not None
        lines = content.decode("utf-8-sig").split("\n")

        # 应该只有表头
        assert len(lines) == 2  # 表头 + 空行

    def test_generate_csv_comma_escape(self):
        """Test CSV properly escapes commas."""
        data = [
            {
                "item_id": "item_001",
                "title": "标题,包含,逗号",
                "viral_type": None,
                "video_url": None,
                "star_id": "star_001",
                "star_unique_id": "unique_001",
                "star_nickname": "测试达人",
                "publish_time": None,
                "natural_play_cnt": 0,
                "heated_play_cnt": 0,
                "total_play_cnt": 0,
                "total_interact": 0,
                "like_cnt": 0,
                "share_cnt": 0,
                "comment_cnt": 0,
                "new_a3_rate": None,
                "after_view_search_uv": 0,
                "return_search_cnt": 0,
                "industry_id": None,
                "industry_name": None,
                "brand_id": None,
                "brand_name": None,
                "estimated_video_cost": 0,
                "estimated_natural_cpm": None,
                "estimated_natural_search_uv": None,
                "estimated_natural_search_cost": None,
            }
        ]
        content = generate_csv(data)
        csv_text = content.decode("utf-8-sig")

        # 包含逗号的字段应该被引号包裹
        assert '"标题,包含,逗号"' in csv_text

    def test_column_headers_complete(self):
        """Test that all required columns are defined."""
        expected_columns = [
            "视频ID",
            "视频标题",
            "爆文类型",
            "视频链接",
            "新增A3率",
            "看后搜人数",
            "回搜次数",
            "自然曝光数",
            "加热曝光数",
            "总曝光数",
            "总互动",
            "点赞",
            "转发",
            "评论",
            "合作行业ID",
            "合作行业",
            "合作品牌ID",
            "合作品牌",
            "发布时间",
            "达人昵称",
            "达人unique_id",
            "预估视频价格",
            "预估自然CPM",
            "预估自然看后搜人数",
            "预估自然看后搜人数成本",
        ]

        for col in expected_columns:
            assert col in [h[0] for h in COLUMN_HEADERS], f"Missing column: {col}"