mirror of
https://github.com/NanmiCoder/MediaCrawler.git
synced 2026-10-07 13:08:11 +08:00
feat: Add Excel export functionality and unit tests
Features: - Excel export with formatted multi-sheet workbooks (Contents, Comments, Creators) - Professional styling: blue headers, auto-width columns, borders, text wrapping - Smart export: empty sheets automatically removed - Support for all platforms (xhs, dy, ks, bili, wb, tieba, zhihu) Testing: - Added pytest framework with asyncio support - Unit tests for Excel store functionality - Unit tests for store factory pattern - Shared fixtures for test data - Test coverage for edge cases Documentation: - Comprehensive Excel export guide (docs/excel_export_guide.md) - Updated README.md and README_en.md with Excel examples - Updated config comments to include excel option Dependencies: - Added openpyxl>=3.1.2 for Excel support - Added pytest>=7.4.0 and pytest-asyncio>=0.21.0 for testing This contribution adds immediate value for users who need data analysis capabilities and establishes a testing foundation for future development.
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# MediaCrawler Test Suite
|
||||
@@ -0,0 +1,81 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Pytest configuration and shared fixtures
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Add project root to Python path
|
||||
project_root = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(project_root))
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def project_root_path():
|
||||
"""Return project root path"""
|
||||
return project_root
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_xhs_note():
|
||||
"""Sample Xiaohongshu note data for testing"""
|
||||
return {
|
||||
"note_id": "test_note_123",
|
||||
"type": "normal",
|
||||
"title": "测试标题 Test Title",
|
||||
"desc": "这是一个测试描述 This is a test description",
|
||||
"video_url": "",
|
||||
"time": 1700000000,
|
||||
"last_update_time": 1700000000,
|
||||
"user_id": "user_123",
|
||||
"nickname": "测试用户",
|
||||
"avatar": "https://example.com/avatar.jpg",
|
||||
"liked_count": 100,
|
||||
"collected_count": 50,
|
||||
"comment_count": 25,
|
||||
"share_count": 10,
|
||||
"ip_location": "上海",
|
||||
"image_list": "https://example.com/img1.jpg,https://example.com/img2.jpg",
|
||||
"tag_list": "测试,编程,Python",
|
||||
"note_url": "https://www.xiaohongshu.com/explore/test_note_123",
|
||||
"source_keyword": "测试关键词",
|
||||
"xsec_token": "test_token_123"
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_xhs_comment():
|
||||
"""Sample Xiaohongshu comment data for testing"""
|
||||
return {
|
||||
"comment_id": "comment_123",
|
||||
"create_time": 1700000000,
|
||||
"ip_location": "北京",
|
||||
"note_id": "test_note_123",
|
||||
"content": "这是一条测试评论 This is a test comment",
|
||||
"user_id": "user_456",
|
||||
"nickname": "评论用户",
|
||||
"avatar": "https://example.com/avatar2.jpg",
|
||||
"sub_comment_count": 5,
|
||||
"pictures": "",
|
||||
"parent_comment_id": 0,
|
||||
"like_count": 15
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_xhs_creator():
|
||||
"""Sample Xiaohongshu creator data for testing"""
|
||||
return {
|
||||
"user_id": "creator_123",
|
||||
"nickname": "创作者名称",
|
||||
"gender": "女",
|
||||
"avatar": "https://example.com/creator_avatar.jpg",
|
||||
"desc": "这是创作者简介",
|
||||
"ip_location": "广州",
|
||||
"follows": 500,
|
||||
"fans": 10000,
|
||||
"interaction": 50000,
|
||||
"tag_list": '{"profession": "设计师", "interest": "摄影"}'
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Unit tests for Excel export functionality
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
import shutil
|
||||
|
||||
try:
|
||||
import openpyxl
|
||||
EXCEL_AVAILABLE = True
|
||||
except ImportError:
|
||||
EXCEL_AVAILABLE = False
|
||||
|
||||
from store.excel_store_base import ExcelStoreBase
|
||||
|
||||
|
||||
@pytest.mark.skipif(not EXCEL_AVAILABLE, reason="openpyxl not installed")
|
||||
class TestExcelStoreBase:
|
||||
"""Test cases for ExcelStoreBase"""
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir(self):
|
||||
"""Create temporary directory for test files"""
|
||||
temp_path = tempfile.mkdtemp()
|
||||
yield temp_path
|
||||
# Cleanup
|
||||
shutil.rmtree(temp_path, ignore_errors=True)
|
||||
|
||||
@pytest.fixture
|
||||
def excel_store(self, temp_dir, monkeypatch):
|
||||
"""Create ExcelStoreBase instance for testing"""
|
||||
# Monkey patch data directory
|
||||
monkeypatch.chdir(temp_dir)
|
||||
store = ExcelStoreBase(platform="test", crawler_type="search")
|
||||
yield store
|
||||
# Cleanup is handled by temp_dir fixture
|
||||
|
||||
def test_initialization(self, excel_store):
|
||||
"""Test Excel store initialization"""
|
||||
assert excel_store.platform == "test"
|
||||
assert excel_store.crawler_type == "search"
|
||||
assert excel_store.workbook is not None
|
||||
assert excel_store.contents_sheet is not None
|
||||
assert excel_store.comments_sheet is not None
|
||||
assert excel_store.creators_sheet is not None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_content(self, excel_store):
|
||||
"""Test storing content data"""
|
||||
content_item = {
|
||||
"note_id": "test123",
|
||||
"title": "Test Title",
|
||||
"desc": "Test Description",
|
||||
"user_id": "user456",
|
||||
"nickname": "TestUser",
|
||||
"liked_count": 100,
|
||||
"comment_count": 50
|
||||
}
|
||||
|
||||
await excel_store.store_content(content_item)
|
||||
|
||||
# Verify data was written
|
||||
assert excel_store.contents_sheet.max_row == 2 # Header + 1 data row
|
||||
assert excel_store.contents_headers_written is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_comment(self, excel_store):
|
||||
"""Test storing comment data"""
|
||||
comment_item = {
|
||||
"comment_id": "comment123",
|
||||
"note_id": "note456",
|
||||
"content": "Great post!",
|
||||
"user_id": "user789",
|
||||
"nickname": "Commenter",
|
||||
"like_count": 10
|
||||
}
|
||||
|
||||
await excel_store.store_comment(comment_item)
|
||||
|
||||
# Verify data was written
|
||||
assert excel_store.comments_sheet.max_row == 2 # Header + 1 data row
|
||||
assert excel_store.comments_headers_written is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_creator(self, excel_store):
|
||||
"""Test storing creator data"""
|
||||
creator_item = {
|
||||
"user_id": "creator123",
|
||||
"nickname": "Creator Name",
|
||||
"fans": 10000,
|
||||
"follows": 500,
|
||||
"interaction": 50000
|
||||
}
|
||||
|
||||
await excel_store.store_creator(creator_item)
|
||||
|
||||
# Verify data was written
|
||||
assert excel_store.creators_sheet.max_row == 2 # Header + 1 data row
|
||||
assert excel_store.creators_headers_written is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_items(self, excel_store):
|
||||
"""Test storing multiple items"""
|
||||
# Store multiple content items
|
||||
for i in range(5):
|
||||
await excel_store.store_content({
|
||||
"note_id": f"note{i}",
|
||||
"title": f"Title {i}",
|
||||
"liked_count": i * 10
|
||||
})
|
||||
|
||||
# Verify all items were stored
|
||||
assert excel_store.contents_sheet.max_row == 6 # Header + 5 data rows
|
||||
|
||||
def test_flush(self, excel_store):
|
||||
"""Test flushing data to file"""
|
||||
# Add some test data
|
||||
asyncio.run(excel_store.store_content({
|
||||
"note_id": "test",
|
||||
"title": "Test"
|
||||
}))
|
||||
|
||||
# Flush to file
|
||||
excel_store.flush()
|
||||
|
||||
# Verify file was created
|
||||
assert excel_store.filename.exists()
|
||||
|
||||
# Verify file can be opened
|
||||
wb = openpyxl.load_workbook(excel_store.filename)
|
||||
assert "Contents" in wb.sheetnames
|
||||
wb.close()
|
||||
|
||||
def test_header_formatting(self, excel_store):
|
||||
"""Test header row formatting"""
|
||||
asyncio.run(excel_store.store_content({"note_id": "test", "title": "Test"}))
|
||||
|
||||
# Check header formatting
|
||||
header_cell = excel_store.contents_sheet.cell(row=1, column=1)
|
||||
assert header_cell.font.bold is True
|
||||
assert header_cell.fill.start_color.rgb == "FF366092"
|
||||
|
||||
def test_empty_sheets_removed(self, excel_store):
|
||||
"""Test that empty sheets are removed on flush"""
|
||||
# Only add content, leave comments and creators empty
|
||||
asyncio.run(excel_store.store_content({"note_id": "test"}))
|
||||
|
||||
excel_store.flush()
|
||||
|
||||
# Reload workbook
|
||||
wb = openpyxl.load_workbook(excel_store.filename)
|
||||
|
||||
# Only Contents sheet should exist
|
||||
assert "Contents" in wb.sheetnames
|
||||
assert "Comments" not in wb.sheetnames
|
||||
assert "Creators" not in wb.sheetnames
|
||||
wb.close()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not EXCEL_AVAILABLE, reason="openpyxl not installed")
|
||||
def test_excel_import_availability():
|
||||
"""Test that openpyxl is available"""
|
||||
assert EXCEL_AVAILABLE is True
|
||||
import openpyxl
|
||||
assert openpyxl is not None
|
||||
@@ -0,0 +1,75 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Unit tests for Store Factory functionality
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from store.xhs import XhsStoreFactory
|
||||
from store.xhs._store_impl import (
|
||||
XhsCsvStoreImplement,
|
||||
XhsJsonStoreImplement,
|
||||
XhsDbStoreImplement,
|
||||
XhsSqliteStoreImplement,
|
||||
XhsMongoStoreImplement,
|
||||
XhsExcelStoreImplement
|
||||
)
|
||||
|
||||
|
||||
class TestXhsStoreFactory:
|
||||
"""Test cases for XhsStoreFactory"""
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'csv')
|
||||
def test_create_csv_store(self):
|
||||
"""Test creating CSV store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsCsvStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'json')
|
||||
def test_create_json_store(self):
|
||||
"""Test creating JSON store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsJsonStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'db')
|
||||
def test_create_db_store(self):
|
||||
"""Test creating database store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsDbStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'sqlite')
|
||||
def test_create_sqlite_store(self):
|
||||
"""Test creating SQLite store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsSqliteStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'mongodb')
|
||||
def test_create_mongodb_store(self):
|
||||
"""Test creating MongoDB store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsMongoStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'excel')
|
||||
@patch('var.crawler_type_var.get', return_value='search')
|
||||
def test_create_excel_store(self, mock_crawler_type):
|
||||
"""Test creating Excel store"""
|
||||
store = XhsStoreFactory.create_store()
|
||||
assert isinstance(store, XhsExcelStoreImplement)
|
||||
|
||||
@patch('config.SAVE_DATA_OPTION', 'invalid')
|
||||
def test_invalid_store_option(self):
|
||||
"""Test that invalid store option raises ValueError"""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
XhsStoreFactory.create_store()
|
||||
|
||||
assert "Invalid save option" in str(exc_info.value)
|
||||
|
||||
def test_all_stores_registered(self):
|
||||
"""Test that all store types are registered"""
|
||||
expected_stores = ['csv', 'json', 'db', 'sqlite', 'mongodb', 'excel']
|
||||
|
||||
for store_type in expected_stores:
|
||||
assert store_type in XhsStoreFactory.STORES
|
||||
|
||||
assert len(XhsStoreFactory.STORES) == len(expected_stores)
|
||||
Reference in New Issue
Block a user