Files
drive/tests/test_storage_cleanup.py
T
2026-10-18 23:52:30 +05:30

135 lines
5.1 KiB
Python

"""Temporary-storage safety: cleanup never touches anything outside the app's temp files."""
import os
from pathlib import Path
import pytest
from caption_parser import parse_caption
from config import ConfigError
from conftest import CHANNEL_ID, FakeDrive, FakeSource, build_queue, make_config, run, video_bytes
from downloader import Downloader
from models import State
PREFIX = abs(CHANNEL_ID)
def _record(db, source, message_id, state):
msg = source.add(message_id, f"Title: V{message_id}.mp4\nTopic: T")
db.insert_video(msg, parse_caption(msg.caption))
path = {
State.COMPLETED: [State.DOWNLOADING, State.DOWNLOADED, State.UPLOADING, State.COMPLETED],
State.FAILED: [State.FAILED],
State.SKIPPED: [State.SKIPPED],
State.DOWNLOADED: [State.DOWNLOADING, State.DOWNLOADED],
State.UPLOADING: [State.DOWNLOADING, State.DOWNLOADED, State.UPLOADING],
State.PENDING: [],
}[state]
for step in path:
db.transition((CHANNEL_ID, message_id), step)
def test_cleanup_only_removes_finished_app_temp_files(config, db, source, tmp_path):
downloader = Downloader(source, config.download_dir)
d = config.download_dir
for message_id, state in [(1, State.COMPLETED), (2, State.SKIPPED), (3, State.FAILED), (4, State.PENDING),
(5, State.DOWNLOADED), (6, State.UPLOADING)]:
_record(db, source, message_id, state)
files = {
"done": d / f"{PREFIX}_1.video",
"skipped": d / f"{PREFIX}_2.part",
"failed": d / f"{PREFIX}_3.video",
"pending_partial": d / f"{PREFIX}_4.part",
"downloaded": d / f"{PREFIX}_5.video",
"downloaded_stale_part": d / f"{PREFIX}_5.part",
"uploading": d / f"{PREFIX}_6.video",
"unknown_record": d / f"{PREFIX}_99.video",
"user_file": d / "notes.txt",
"lookalike": d / f"{PREFIX}_1.video.bak",
}
for path in files.values():
path.write_bytes(b"x")
outside = tmp_path / "outside" / f"{PREFIX}_1.video"
outside.parent.mkdir()
outside.write_bytes(b"precious")
(d / f"{PREFIX}_7.video").symlink_to(outside) # matching name, but a symlink pointing outside
subdir = d / f"{PREFIX}_8.video"
subdir.mkdir() # directories are never removed
deleted = downloader.cleanup_orphans(db)
gone = {k for k, p in files.items() if not p.exists()}
assert gone == {"done", "skipped", "failed", "downloaded_stale_part", "unknown_record"}
assert deleted == 5
assert outside.read_bytes() == b"precious"
assert (d / f"{PREFIX}_7.video").is_symlink() and subdir.is_dir()
def test_failed_files_can_be_kept(config, db, source):
downloader = Downloader(source, config.download_dir)
_record(db, source, 3, State.FAILED)
path = config.download_dir / f"{PREFIX}_3.video"
path.write_bytes(b"x")
assert downloader.cleanup_orphans(db, keep_failed=True) == 0 and path.exists()
def test_safe_unlink_refuses_anything_outside_temp_dir(config, source, tmp_path):
downloader = Downloader(source, config.download_dir)
outside = tmp_path / f"{PREFIX}_1.video"
outside.write_bytes(b"keep me")
assert downloader.safe_unlink(outside) is False and outside.exists()
traversal = config.download_dir / ".." / f"{PREFIX}_1.video"
assert downloader.safe_unlink(traversal) is False and outside.exists()
odd = config.download_dir / "important.db"
odd.write_bytes(b"x")
assert downloader.safe_unlink(odd) is False and odd.exists()
def test_low_disk_space_reclaims_abandoned_files_before_failing(config, db, source):
drive = FakeDrive()
_record(db, source, 1, State.COMPLETED)
orphan = config.download_dir
orphan.mkdir(parents=True, exist_ok=True)
orphan = orphan / f"{PREFIX}_1.video"
orphan.write_bytes(b"x" * 10)
source.add(2, "Title: New.mp4\nTopic: T", video_bytes(600_000, 2))
def disk_usage(_):
free = 10**12 if not orphan.exists() else 0 # only enough space once the orphan is gone
return (10**12, 10**12 - free, free)
queue = build_queue(config, db, source, drive, disk_usage=disk_usage)
async def go():
await queue.sync()
queue.start()
await queue.wait()
run(go())
assert not orphan.exists()
assert db.get((CHANNEL_ID, 2)).state == State.COMPLETED
def test_permanent_download_failure_removes_partial(config, db, source):
from errors import PermanentError
drive = FakeDrive()
source.add(1, "Title: A.mp4\nTopic: T", video_bytes(3 * 1024 * 1024, 1))
source.failures[1] = [("after", 2 * 1024 * 1024, PermanentError("message deleted mid-download"))]
queue = build_queue(config, db, source, drive)
async def go():
await queue.sync()
queue.start()
await queue.wait()
run(go())
assert db.get((CHANNEL_ID, 1)).state == State.FAILED
assert list(config.download_dir.iterdir()) == []
@pytest.mark.parametrize("bad", ["/", str(Path.home()), "/tmp"])
def test_download_dir_must_be_dedicated(tmp_path, bad):
with pytest.raises(ConfigError) as err:
make_config(tmp_path, DOWNLOAD_DIR=bad)
assert "DOWNLOAD_DIR" in str(err.value)