diff --git a/src/wordstat/collector.py b/src/wordstat/collector.py index 3f7d755..cb23f9c 100644 --- a/src/wordstat/collector.py +++ b/src/wordstat/collector.py @@ -7,6 +7,7 @@ import tempfile import time from collections.abc import Callable +from dataclasses import dataclass from datetime import UTC, date, datetime, timedelta from pathlib import Path @@ -17,6 +18,7 @@ from wordstat.errors import ( AuthenticationRequiredError, DownloadEscapedError, + DownloadNoNewPathError, DownloadTimeoutError, InterfaceChangedError, InvalidRequestError, @@ -166,6 +168,14 @@ def _without_traceback(error: Exception) -> Exception: return error.with_traceback(None) +@dataclass(frozen=True) +class _RetryExportResult: + """A parsed export whose path and contents are guaranteed to match.""" + + source: Path + dataset: CsvDataset + + class WordstatCollector: """Export four Wordstat reports through an already authenticated CDP session.""" @@ -589,12 +599,16 @@ async def _collect_one( # separate from the fail-closed predicate below. dataset = parse_wordstat_csv(source, view) if _should_retry_empty_export(view, dataset) or _is_untrustworthy_empty_export(view, dataset): - if self.empty_export_retry_seconds > 0: - await asyncio.sleep(self.empty_export_retry_seconds) - source, escape_warning = await self._download_current_view(page, session, downloads_path) - if escape_warning is not None: - escaped_download_warnings.append(f"[{view.value}] {escape_warning}") - dataset = parse_wordstat_csv(source, view) + retry = await self._retry_empty_export( + page, + session, + downloads_path, + source, + dataset, + view, + escaped_download_warnings, + ) + source, dataset = retry.source, retry.dataset if _is_untrustworthy_empty_export(view, dataset): raise InterfaceChangedError( f"Wordstat returned an empty {view.value} CSV after a retry, but the page had " @@ -1181,6 +1195,55 @@ async def _table_snapshot(self, page) -> str | None: f"() => document.querySelector({json.dumps(TABLE_ROW_SELECTOR)})?.textContent ?? null" ) + async def _retry_empty_export( + self, + page, + session: BrowserSession, + downloads_path: Path, + source: Path, + dataset: CsvDataset, + view: WordstatView, + escaped_download_warnings: list[str], + ) -> _RetryExportResult: + """Retry an empty export and return a matching path/dataset pair. + + A no-new-path timeout can mean Chrome reused and overwrote ``source``. + The original is backed up before the retry; on that specific signal, + restore it to the original path before returning. All other failures + propagate unchanged. The single outer ``finally`` owns the temporary + file through creation, copy, retry, restore, and cleanup. + """ + if self.empty_export_retry_seconds > 0: + await asyncio.sleep(self.empty_export_retry_seconds) + retry_backup: Path | None = None + try: + with tempfile.NamedTemporaryFile( + prefix=f".{view.value}-retry-", + suffix=".csv", + dir=self.output_root, + delete=False, + ) as backup_file: + retry_backup = Path(backup_file.name) + shutil.copy2(source, retry_backup) + try: + retry_source, retry_escape_warning = await self._download_current_view( + page, session, downloads_path + ) + except DownloadNoNewPathError: + if not _should_retry_empty_export(view, dataset): + raise + shutil.copy2(retry_backup, source) + return _RetryExportResult(source=source, dataset=dataset) + if retry_escape_warning is not None: + escaped_download_warnings.append(f"[{view.value}] {retry_escape_warning}") + return _RetryExportResult( + source=retry_source, + dataset=parse_wordstat_csv(retry_source, view), + ) + finally: + if retry_backup is not None: + retry_backup.unlink(missing_ok=True) + async def _download_current_view( self, page, session: BrowserSession, downloads_path: Path ) -> tuple[Path, str | None]: @@ -1279,7 +1342,7 @@ async def _download_current_view( "automatically. Move it manually if it belongs to this run." ) await asyncio.sleep(0.25) - raise DownloadTimeoutError("Wordstat did not produce a CSV before the download timeout") + raise DownloadNoNewPathError("Wordstat did not produce a new CSV before the download timeout") async def _click(self, page, selector: str) -> None: result = await page.evaluate( diff --git a/src/wordstat/errors.py b/src/wordstat/errors.py index df500ff..c6b38d8 100644 --- a/src/wordstat/errors.py +++ b/src/wordstat/errors.py @@ -29,6 +29,14 @@ class DownloadTimeoutError(WordstatError): """The UI accepted an export request but did not produce a CSV file.""" +class DownloadNoNewPathError(DownloadTimeoutError): + """The export completed without exposing a new file path. + + This is distinct from other download failures (notably multiple new CSVs) + because Chrome can handle a repeated export by reusing the existing path. + """ + + class DownloadEscapedError(WordstatError): """Chrome reported a download outside the run's own downloads directory. diff --git a/tests/test_collector_batch.py b/tests/test_collector_batch.py index 2ef1c1b..7bcf4c0 100644 --- a/tests/test_collector_batch.py +++ b/tests/test_collector_batch.py @@ -9,13 +9,16 @@ import asyncio from datetime import UTC, datetime +import pyarrow.parquet as parquet import pytest import wordstat.collector as collector_module from wordstat.collector import WordstatCollector from wordstat.errors import ( AuthenticationRequiredError, + DownloadNoNewPathError, DownloadTimeoutError, + InterfaceChangedError, InvalidRequestError, PhraseEntryError, ResumeMismatchError, @@ -463,7 +466,9 @@ async def fake_download(self, page, session, dl_path): monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) - collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + collector = WordstatCollector( + "cdp", tmp_path, keep_raw=True, settling_seconds=0, empty_export_retry_seconds=0 + ) result = asyncio.run( collector._collect_one( _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False @@ -473,6 +478,11 @@ async def fake_download(self, page, session, dl_path): assert download_count == 6 # top views each needed one content-based retry assert {export.view for export in result.manifest.exports} == set(WordstatView) assert all(export.row_count == 1 for export in result.manifest.exports) + assert not list(tmp_path.glob("*retry*")) + top_export = next(export for export in result.manifest.exports if export.view is WordstatView.TOP_POPULAR) + assert top_export.raw_file == "top_popular.csv" + assert "январь 2024;100" in (result.run_directory / top_export.raw_file).read_text(encoding="cp1251") + assert parquet.read_table(result.run_directory / top_export.file).num_rows == top_export.row_count def test_collect_one_accepts_persistent_empty_top_exports_after_retry(monkeypatch, tmp_path): @@ -521,6 +531,227 @@ async def fake_download(self, page, session, dl_path): assert all(export.row_count == 1 for export in result.manifest.exports[2:]) +def test_collect_one_keeps_empty_top_export_when_retry_times_out(monkeypatch, tmp_path): + """A failed top-view retry must fall back to the first empty CSV. + + Chrome can overwrite the first download with the same filename, leaving + the snapshot-based downloader with no new path to observe. That timeout + must not prevent the other views from being collected. + """ + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + download_count = 0 + + async def fake_download(self, page, session, dl_path): + nonlocal download_count + download_count += 1 + if download_count == 1: + source = dl_path / "export.csv" + _write_empty_view_csv(source) + return source, None + if download_count == 2: + raise DownloadNoNewPathError("simulated same-name retry timeout") + source = dl_path / f"export-{download_count}.csv" + _write_view_csv(source, "тест") + return source, None + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + result = asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + assert download_count == 5 + assert result.view_errors == {} + assert [export.row_count for export in result.manifest.exports] == [0, 1, 1, 1] + assert result.manifest.empty_views == [WordstatView.TOP_POPULAR] + + +def test_collect_one_preserves_first_csv_when_retry_overwrites_path(monkeypatch, tmp_path): + """Fallback parquet and raw CSV must describe the same first export.""" + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + download_count = 0 + + async def fake_download(self, page, session, dl_path): + nonlocal download_count + download_count += 1 + source = dl_path / "export.csv" + if download_count == 1: + _write_empty_view_csv(source) + return source, None + if download_count == 2: + _write_view_csv(source, "тест") + raise DownloadNoNewPathError("simulated overwritten-path timeout") + source = dl_path / f"export-{download_count}.csv" + _write_view_csv(source, "тест") + return source, None + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector( + "cdp", tmp_path, keep_raw=True, settling_seconds=0, empty_export_retry_seconds=0 + ) + result = asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + assert result.manifest.exports[0].row_count == 0 + assert (result.run_directory / "top_popular.csv").read_text(encoding="cp1251") == "Запрос;Показов\n" + + +def test_collect_one_cleans_retry_backup_when_copy_fails(monkeypatch, tmp_path): + """A failed backup copy must not leave a temporary file behind.""" + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + async def fake_download(self, page, session, dl_path): + source = dl_path / "export.csv" + _write_empty_view_csv(source) + return source, None + + def failing_copy(source, destination): + raise OSError("simulated backup filesystem failure") + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + monkeypatch.setattr(collector_module.shutil, "copy2", failing_copy) + + collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + with pytest.raises(OSError, match="backup filesystem failure"): + asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + assert not list(tmp_path.glob("*retry*")) + + +def test_collect_one_does_not_swallow_non_timeout_top_retry_error(monkeypatch, tmp_path): + """Only the known same-name timeout is recoverable for a top export.""" + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + async def fake_download(self, page, session, dl_path): + source = dl_path / "export.csv" + if source.exists(): + raise InterfaceChangedError("simulated changed export control") + _write_empty_view_csv(source) + return source, None + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + with pytest.raises(InterfaceChangedError, match="changed export control"): + asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + +def test_collect_one_does_not_swallow_ambiguous_top_retry_timeout(monkeypatch, tmp_path): + """A multi-download timeout is not the recoverable no-new-path signal.""" + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + async def fake_download(self, page, session, dl_path): + source = dl_path / "export.csv" + if source.exists(): + raise DownloadTimeoutError("Wordstat produced more than one new CSV for a single export") + _write_empty_view_csv(source) + return source, None + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + with pytest.raises(DownloadTimeoutError, match="more than one new CSV"): + asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + +def test_collect_one_does_not_swallow_dynamics_retry_timeout(monkeypatch, tmp_path): + """An empty dynamics retry timeout must preserve fail-closed behavior.""" + + _patch_common(monkeypatch) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + download_count = 0 + + async def fake_download(self, page, session, dl_path): + nonlocal download_count + download_count += 1 + source = dl_path / f"export-{download_count}.csv" + if download_count == 3: + _write_empty_view_csv(source) + return source, None + if download_count == 4: + raise DownloadNoNewPathError("simulated dynamics retry timeout") + _write_view_csv(source, "тест") + return source, None + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path, settling_seconds=0, empty_export_retry_seconds=0) + result = asyncio.run( + collector._collect_one( + _FakePage(), _FakeSession(), downloads_path, "тест", "Россия", set_region=False + ) + ) + + assert result.manifest.missing_views == [WordstatView.DYNAMICS, WordstatView.REGIONS] + assert result.manifest.exports[0].view is WordstatView.TOP_POPULAR + assert result.manifest.exports[1].view is WordstatView.TOP_RELATED + assert result.view_errors[WordstatView.DYNAMICS] == ( + "DownloadNoNewPathError: simulated dynamics retry timeout" + ) + + def test_collect_one_keeps_dynamics_empty_export_fail_closed(monkeypatch, tmp_path): """The top-view retry must not change DYNAMICS' existing safety path."""