test_resume_tracker.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336
  1. """
  2. Unit tests for the resume tracker module.
  3. Tests cover resume state tracking and loading.
  4. """
  5. import sys
  6. from unittest.mock import Mock
  7. from pathlib import Path
  8. import json
  9. import tempfile
  10. # Mock torch and transformers before importing
  11. sys_mock = Mock()
  12. sys.modules["torch"] = sys_mock
  13. sys.modules["transformers"] = sys_mock
  14. import pytest
  15. from src.translator.resume_tracker import ResumeTracker, ResumeState
  16. class TestResumeState:
  17. """Test cases for ResumeState dataclass."""
  18. def test_create_resume_state(self):
  19. """Test creating a resume state."""
  20. state = ResumeState(
  21. work_id="test_work",
  22. last_completed_index=10,
  23. last_failed_index=None,
  24. total_chapters=100,
  25. )
  26. assert state.work_id == "test_work"
  27. assert state.last_completed_index == 10
  28. assert state.last_failed_index is None
  29. assert state.total_chapters == 100
  30. def test_to_dict(self):
  31. """Test converting resume state to dictionary."""
  32. state = ResumeState(
  33. work_id="test_work",
  34. last_completed_index=10,
  35. last_failed_index=None,
  36. total_chapters=100,
  37. resume_message="Test message",
  38. )
  39. data = state.to_dict()
  40. assert data["work_id"] == "test_work"
  41. assert data["last_completed_index"] == 10
  42. assert data["last_failed_index"] is None
  43. assert data["total_chapters"] == 100
  44. assert data["resume_message"] == "Test message"
  45. assert "timestamp" in data
  46. def test_from_dict(self):
  47. """Test creating resume state from dictionary."""
  48. data = {
  49. "work_id": "test_work",
  50. "last_completed_index": 10,
  51. "last_failed_index": None,
  52. "total_chapters": 100,
  53. "timestamp": "2024-01-01T00:00:00",
  54. "resume_message": "Test message",
  55. }
  56. state = ResumeState.from_dict(data)
  57. assert state.work_id == "test_work"
  58. assert state.last_completed_index == 10
  59. assert state.total_chapters == 100
  60. def test_get_resume_index_with_failure(self):
  61. """Test getting resume index when there's a failed chapter."""
  62. state = ResumeState(
  63. work_id="test_work",
  64. last_completed_index=10,
  65. last_failed_index=15,
  66. total_chapters=100,
  67. )
  68. assert state.get_resume_index() == 15
  69. def test_get_resume_index_without_failure(self):
  70. """Test getting resume index when there's no failed chapter."""
  71. state = ResumeState(
  72. work_id="test_work",
  73. last_completed_index=10,
  74. last_failed_index=None,
  75. total_chapters=100,
  76. )
  77. assert state.get_resume_index() == 11
  78. def test_is_resumable_with_remaining(self):
  79. """Test is_resumable when there are remaining chapters."""
  80. state = ResumeState(
  81. work_id="test_work",
  82. last_completed_index=10,
  83. last_failed_index=None,
  84. total_chapters=100,
  85. )
  86. assert state.is_resumable() is True
  87. def test_is_resumable_complete(self):
  88. """Test is_resumable when all chapters are complete."""
  89. state = ResumeState(
  90. work_id="test_work",
  91. last_completed_index=99,
  92. last_failed_index=None,
  93. total_chapters=100,
  94. )
  95. assert state.is_resumable() is False
  96. class TestResumeTracker:
  97. """Test cases for ResumeTracker class."""
  98. @pytest.fixture
  99. def temp_dir(self):
  100. """Create a temporary directory for testing."""
  101. with tempfile.TemporaryDirectory() as tmp:
  102. yield Path(tmp)
  103. @pytest.fixture
  104. def tracker(self, temp_dir):
  105. """Create a resume tracker for testing."""
  106. return ResumeTracker(temp_dir)
  107. def test_init(self, temp_dir):
  108. """Test ResumeTracker initialization."""
  109. tracker = ResumeTracker(temp_dir)
  110. assert tracker.storage_dir == temp_dir
  111. assert tracker.resume_dir == temp_dir / "resume"
  112. assert tracker.resume_dir.exists()
  113. def test_save_resume_state(self, tracker):
  114. """Test saving resume state."""
  115. state = tracker.save_resume_state(
  116. work_id="test_work",
  117. last_completed_index=10,
  118. last_failed_index=None,
  119. total_chapters=100,
  120. )
  121. assert state.work_id == "test_work"
  122. assert state.last_completed_index == 10
  123. assert "Resume from chapter" in state.resume_message
  124. # Check file was created
  125. resume_file = tracker._get_resume_file_path("test_work")
  126. assert resume_file.exists()
  127. def test_load_resume_state(self, tracker):
  128. """Test loading resume state."""
  129. # Save state first
  130. tracker.save_resume_state(
  131. work_id="test_work",
  132. last_completed_index=10,
  133. last_failed_index=None,
  134. total_chapters=100,
  135. )
  136. # Load it back
  137. loaded = tracker.load_resume_state("test_work")
  138. assert loaded is not None
  139. assert loaded.work_id == "test_work"
  140. assert loaded.last_completed_index == 10
  141. assert loaded.total_chapters == 100
  142. def test_load_nonexistent_resume_state(self, tracker):
  143. """Test loading a non-existent resume state."""
  144. loaded = tracker.load_resume_state("nonexistent")
  145. assert loaded is None
  146. def test_delete_resume_state(self, tracker):
  147. """Test deleting resume state."""
  148. # Save state first
  149. tracker.save_resume_state(
  150. work_id="test_work",
  151. last_completed_index=10,
  152. last_failed_index=None,
  153. total_chapters=100,
  154. )
  155. # Delete it
  156. tracker.delete_resume_state("test_work")
  157. # Check it's gone
  158. resume_file = tracker._get_resume_file_path("test_work")
  159. assert not resume_file.exists()
  160. def test_get_resume_index(self, tracker):
  161. """Test getting resume index."""
  162. # Save state first
  163. tracker.save_resume_state(
  164. work_id="test_work",
  165. last_completed_index=10,
  166. last_failed_index=None,
  167. total_chapters=100,
  168. )
  169. # Get resume index
  170. index = tracker.get_resume_index("test_work", 100)
  171. assert index == 11
  172. def test_get_resume_index_nonexistent(self, tracker):
  173. """Test getting resume index for non-existent state."""
  174. index = tracker.get_resume_index("nonexistent", 100)
  175. assert index is None
  176. def test_update_on_chapter_complete(self, tracker):
  177. """Test updating state on chapter completion."""
  178. tracker.update_on_chapter_complete(
  179. work_id="test_work",
  180. chapter_index=5,
  181. total_chapters=100,
  182. )
  183. loaded = tracker.load_resume_state("test_work")
  184. assert loaded is not None
  185. assert loaded.last_completed_index == 5
  186. assert loaded.last_failed_index is None
  187. def test_update_on_chapter_failed(self, tracker):
  188. """Test updating state on chapter failure."""
  189. tracker.update_on_chapter_failed(
  190. work_id="test_work",
  191. chapter_index=5,
  192. total_chapters=100,
  193. )
  194. loaded = tracker.load_resume_state("test_work")
  195. assert loaded is not None
  196. assert loaded.last_failed_index == 5
  197. def test_generate_resume_message_with_failure(self):
  198. """Test generating resume message with failed chapter."""
  199. state = ResumeState(
  200. work_id="test_work",
  201. last_completed_index=10,
  202. last_failed_index=15,
  203. total_chapters=100,
  204. )
  205. tracker = ResumeTracker(Path("/tmp"))
  206. message = tracker._generate_resume_message(state)
  207. assert "failed chapter 15" in message
  208. assert "85" in message # remaining chapters
  209. def test_generate_resume_message_without_failure(self):
  210. """Test generating resume message without failed chapter."""
  211. state = ResumeState(
  212. work_id="test_work",
  213. last_completed_index=10,
  214. last_failed_index=None,
  215. total_chapters=100,
  216. )
  217. tracker = ResumeTracker(Path("/tmp"))
  218. message = tracker._generate_resume_message(state)
  219. assert "Resume from chapter 11" in message
  220. assert "89" in message # remaining chapters
  221. def test_get_all_resume_states(self, tracker):
  222. """Test getting all resume states."""
  223. # Save multiple states
  224. tracker.save_resume_state("work1", 10, None, 100)
  225. tracker.save_resume_state("work2", 20, None, 200)
  226. states = tracker.get_all_resume_states()
  227. assert len(states) == 2
  228. work_ids = {s.work_id for s in states}
  229. assert work_ids == {"work1", "work2"}
  230. def test_clear_all_resume_states(self, tracker):
  231. """Test clearing all resume states."""
  232. # Save multiple states
  233. tracker.save_resume_state("work1", 10, None, 100)
  234. tracker.save_resume_state("work2", 20, None, 200)
  235. tracker.clear_all_resume_states()
  236. states = tracker.get_all_resume_states()
  237. assert len(states) == 0
  238. def test_update_existing_state_on_complete(self, tracker):
  239. """Test that updating state preserves existing data."""
  240. # Save initial state
  241. tracker.save_resume_state(
  242. work_id="test_work",
  243. last_completed_index=5,
  244. last_failed_index=7,
  245. total_chapters=100,
  246. )
  247. # Update on completion
  248. tracker.update_on_chapter_complete(
  249. work_id="test_work",
  250. chapter_index=10,
  251. total_chapters=100,
  252. )
  253. loaded = tracker.load_resume_state("test_work")
  254. assert loaded.last_completed_index == 10
  255. assert loaded.last_failed_index is None # Should be cleared
  256. def test_custom_resume_message(self, tracker):
  257. """Test saving state with custom resume message."""
  258. state = tracker.save_resume_state(
  259. work_id="test_work",
  260. last_completed_index=10,
  261. last_failed_index=None,
  262. total_chapters=100,
  263. resume_message="Custom message here",
  264. )
  265. assert state.resume_message == "Custom message here"
  266. def test_get_resume_file_path(self, tracker):
  267. """Test getting resume file path."""
  268. path = tracker._get_resume_file_path("test_work")
  269. assert path.name == "test_work_resume.json"
  270. assert path.parent == tracker.resume_dir