返回 douyin-downloader
test_file_manager.py
根目录 / tests / test_file_manager.py
1 from unittest.mock import AsyncMock, MagicMock
2
3 import pytest
4
5 from storage.file_manager import FileManager
6
7
8 def test_file_exists_returns_false_for_missing(tmp_path):
9 fm = FileManager(str(tmp_path))
10 assert fm.file_exists(tmp_path / "nope.mp4") is False
11
12
13 def test_file_exists_returns_false_for_empty(tmp_path):
14 empty = tmp_path / "empty.mp4"
15 empty.write_bytes(b"")
16 fm = FileManager(str(tmp_path))
17 assert fm.file_exists(empty) is False
18
19
20 def test_file_exists_returns_true_for_non_empty(tmp_path):
21 real = tmp_path / "real.mp4"
22 real.write_bytes(b"data")
23 fm = FileManager(str(tmp_path))
24 assert fm.file_exists(real) is True
25
26
27 def test_get_file_size_returns_0_for_missing(tmp_path):
28 fm = FileManager(str(tmp_path))
29 assert fm.get_file_size(tmp_path / "nope.mp4") == 0
30
31
32 def test_get_save_path_creates_directories(tmp_path):
33 fm = FileManager(str(tmp_path))
34 path = fm.get_save_path(
35 "Author", mode="post", aweme_title="Title", aweme_id="123", download_date="2024-01-01"
36 )
37 assert path.exists()
38 assert "Author" in str(path)
39 assert "post" in str(path)
40 assert "123" in str(path)
41
42
43 @pytest.mark.asyncio
44 async def test_download_file_atomic_write(tmp_path):
45 """Downloaded file should appear only after successful completion (atomic rename)."""
46 fm = FileManager(str(tmp_path))
47 save_path = tmp_path / "video.mp4"
48 content = b"fake video content"
49
50 mock_response = AsyncMock()
51 mock_response.status = 200
52 mock_response.content_length = len(content)
53
54 async def iter_chunked(size):
55 yield content
56
57 mock_response.content = MagicMock()
58 mock_response.content.iter_chunked = iter_chunked
59
60 ctx = AsyncMock()
61 ctx.__aenter__ = AsyncMock(return_value=mock_response)
62 ctx.__aexit__ = AsyncMock(return_value=False)
63
64 mock_session = MagicMock()
65 mock_session.get.return_value = ctx
66
67 result = await fm.download_file("https://example.com/v.mp4", save_path, session=mock_session)
68 assert result is True
69 assert save_path.exists()
70 assert save_path.read_bytes() == content
71 assert not save_path.with_suffix(".mp4.tmp").exists()
72
73
74 @pytest.mark.asyncio
75 async def test_download_file_size_mismatch_cleans_up(tmp_path):
76 fm = FileManager(str(tmp_path))
77 save_path = tmp_path / "video.mp4"
78
79 mock_response = AsyncMock()
80 mock_response.status = 200
81 mock_response.content_length = 999
82
83 async def iter_chunked(size):
84 yield b"short"
85
86 mock_response.content = MagicMock()
87 mock_response.content.iter_chunked = iter_chunked
88
89 ctx = AsyncMock()
90 ctx.__aenter__ = AsyncMock(return_value=mock_response)
91 ctx.__aexit__ = AsyncMock(return_value=False)
92
93 mock_session = MagicMock()
94 mock_session.get.return_value = ctx
95
96 result = await fm.download_file("https://example.com/v.mp4", save_path, session=mock_session)
97 assert result is False
98 assert not save_path.exists()
99 assert not save_path.with_suffix(".mp4.tmp").exists()
100
100 lines PYTHON