90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
from .models import FileState
|
|
|
|
|
|
@dataclass
|
|
class FakeFile:
|
|
torrent_hash: str
|
|
torrent_name: str
|
|
file_index: int
|
|
name: str
|
|
size: int
|
|
downloaded: int
|
|
priority: int = 0
|
|
save_path: Path = Path("/downloads")
|
|
|
|
def to_state(self) -> FileState:
|
|
downloaded = min(self.downloaded, self.size)
|
|
remaining = max(self.size - downloaded, 0)
|
|
progress = downloaded / self.size if self.size else 1.0
|
|
return FileState(
|
|
torrent_hash=self.torrent_hash,
|
|
torrent_name=self.torrent_name,
|
|
file_index=self.file_index,
|
|
name=self.name,
|
|
size=self.size,
|
|
downloaded=downloaded,
|
|
progress=progress,
|
|
priority=self.priority,
|
|
remaining=remaining,
|
|
save_path=self.save_path,
|
|
)
|
|
|
|
def clone(self) -> FakeFile:
|
|
return FakeFile(
|
|
torrent_hash=self.torrent_hash,
|
|
torrent_name=self.torrent_name,
|
|
file_index=self.file_index,
|
|
name=self.name,
|
|
size=self.size,
|
|
downloaded=self.downloaded,
|
|
priority=self.priority,
|
|
save_path=self.save_path,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class FakeTorrent:
|
|
torrent_hash: str
|
|
name: str
|
|
files: List[FakeFile]
|
|
sequential: bool = True
|
|
|
|
|
|
class FakeQBClient:
|
|
"""In-memory emulator that mimics the parts of qBittorrent the scheduler needs."""
|
|
|
|
def __init__(self, torrents: List[FakeTorrent]):
|
|
self.torrents: Dict[str, FakeTorrent] = {t.torrent_hash: t for t in torrents}
|
|
self.priority_calls: List[Dict] = []
|
|
self.sequential_calls: List[Dict] = []
|
|
|
|
def fetch_file_states(self) -> List[FileState]:
|
|
states: List[FileState] = []
|
|
for torrent in self.torrents.values():
|
|
for file in torrent.files:
|
|
states.append(file.to_state())
|
|
return states
|
|
|
|
def set_priority(self, torrent_hash: str, file_ids: List[int], priority: int) -> None:
|
|
torrent = self.torrents[torrent_hash]
|
|
for file_id in file_ids:
|
|
for file in torrent.files:
|
|
if file.file_index == file_id:
|
|
file.priority = priority
|
|
break
|
|
self.priority_calls.append(
|
|
{"hash": torrent_hash, "file_ids": list(file_ids), "priority": priority}
|
|
)
|
|
|
|
def set_sequential(self, torrent_hashes: List[str], value: bool) -> None:
|
|
for torrent_hash in torrent_hashes:
|
|
torrent = self.torrents[torrent_hash]
|
|
torrent.sequential = value
|
|
self.sequential_calls.append({"hashes": list(torrent_hashes), "value": value})
|