diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..0286f3c --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,15 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: astral-sh/setup-uv@v3 + - name: ruff check + run: uvx ruff check . diff --git a/CLAUDE.md b/CLAUDE.md index 38d38d7..6fdeff2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -39,6 +39,11 @@ wordstat collect "ремонт квартир" "натяжные потолки" # --cdp-url / WORDSTAT_CDP_URL, --timeout (сек, по умолчанию 45) # --keep-raw — оставить скачанные CSV как .csv рядом с parquet +# дозаписать недостающие виды в уже существующий run-каталог (ровно одна фраза, +# фраза/регион должны совпасть с manifest.json в этом каталоге) +wordstat collect "ремонт квартир" --region "Москва" \ + --resume-dir ./wordstat-output/runs/20260821T090000Z-ремонт-квартир + # тесты pytest pytest tests/test_csv_io.py::test_parse_wordstat_csv_rejects_duplicate_headers # один тест @@ -121,6 +126,52 @@ ruff check . инвариант держится и для сбоя `write_dataset`/`finalize_raw`) скачанный CSV переносится из временного каталога загрузок в `run_directory` фразы и остаётся на диске для разбора, прежде чем исключение уйдёт выше. + - **Манифест пишется инкрементально**, а не одним куском в конце: один раз + до цикла видов (`exports=[]`, честно отражает «прогон начался, ничего ещё + не собрано») и заново после каждого успешно собранного вида + (`merge_export` + `write_manifest`, обе — `storage.py`). Так обрыв + посреди фразы (Ctrl-C, упавший CDP, сбой `write_dataset` на третьем виде) + оставляет на диске манифест, который честно описывает, что реально + собрано, а не либо ничего, либо (что опаснее) манифест с четырьмя видами, + из которых реально записались не все. `status`/`missing_views` в + `CollectionManifest` — `computed_field`, выводятся из `exports`, а не + хранятся отдельно: сконструировать манифест, где они противоречат + `exports`, невозможно (детали — в докстринге модели, `models.py`). + `source_url` при первом прогоне снимается один раз, до цикла видов + (первая вкладка), а не после последнего вида (карта), как было в + однократной записи — это осознанная смена: URL из середины/конца цикла + было бы либо недоступно при инкрементальной записи первого манифеста + (ещё ни одна вкладка не выбрана), либо непостоянно от вида к виду. При + дозаписи `source_url` обновляется заново (см. следующий пункт) — + значение из первого прогона не остаётся висеть навсегда. + - **Дозапись** (`--resume-dir`) — отдельный явный путь, не ослабление + `create_run_directory`: тот каталог создаётся заново, как и раньше; + резюм принимает уже существующий каталог, где `manifest.json` хранит + `phrase`/`region`, и `prepare_resume_directory` (`storage.py`) жёстко + отклоняет (`ResumeMismatchError`) каталог с несовпадающей (после + `.strip()`) фразой или регионом, отсутствующим/битым манифестом или + путём, который не является каталогом — иначе опечатка в `--resume-dir` + молча смешала бы данные двух разных фраз в одном каталоге. Вид считается + уже собранным (`views_to_collect`), только если он одновременно есть в + `manifest.exports` **и** его `.parquet` реально лежит на диске — + удалённый вручную parquet при живой записи в манифесте не должен + восприниматься как «уже собрано». `collect_many`/CLI отклоняют + `--resume-dir` при более чем одной фразе: один run-каталог = одна фраза, + дозапись на батч размазала бы чужие exports по одному манифесту. CLI + дополнительно вызывает `prepare_resume_directory` до старта Chrome + (fail-fast на опечатку), но `_collect_one` всё равно перепроверяет то же + самое перед стартом — CLI-проверка экономит время, а не заменяет гейт. + - **При дозаписи `_collect_one` сразу после `_set_phrase`** (не раньше — + до `_set_phrase` страница ещё показывает предыдущую фразу/вкладку) + перезаписывает манифест с обновлёнными `source_url` и `updated_at`, + оставляя `created_at` нетронутым. Иначе манифест после резюма продолжал + бы утверждать, что все виды собраны в момент первого, прерванного + прогона, и `source_url` указывал бы на устаревшую вкладку/фразу — + ловушка, которую пропустили все три независимых решения issue #2 при + первом проходе. Ранний `return` резюма уже полностью собранного каталога + (нечего досчитывать) этот шаг не выполняет — там реально ничего не + изменилось, бампать `updated_at` было бы ложью. Семантика всех трёх + полей подробно описана в докстринге `CollectionManifest` (`models.py`). - Все клики/проверки идут через `page.evaluate` с CSS-селекторами и строгой проверкой «найден ровно один элемент» — если 0 или >1, кидается `InterfaceChangedError`. Это защита от того, что Wordstat незаметно @@ -193,6 +244,38 @@ ruff check . при `keep_raw` переименовывает в `.csv`. `write_manifest` пишет `manifest.json` в UTF-8 без экранирования кириллицы. После прогона в каталоге лежат четыре `.parquet` и `manifest.json`. + - **`write_manifest` атомарна**: временный файл создаётся в том же + каталоге (`tempfile.mkstemp(dir=path.parent, ...)`), а не в системном + temp — иначе `os.replace` через границу файловой системы падает с + `OSError` вместо атомарного переименования — и заменяет цель через + `os.replace`. Это не деталь реализации, а предпосылка для инкрементальной + записи манифеста (см. `collector.py`): раз манифест теперь переписывается + после каждого вида, а не один раз в конце, окно между усечением файла и + записью нового содержимого стало реальным, и обычный `Path.write_text` + (усечение + запись) оставлял бы пустой/битый `manifest.json` при обрыве + ровно в этот момент. Временный файл всегда убирается в `finally`, даже + если `os.replace` упал — не остаётся мусора вида `.manifest-*.json.tmp`. + Перед `os.replace` — `handle.flush()` + `os.fsync(handle.fileno())`: + `os.replace` гарантирует только порядок операций (переименование), но не + то, что байты временного файла физически дошли до диска — без `fsync` + потеря питания между переименованием и сбросом page cache могла бы + оставить пустой/мусорный `manifest.json` там, где был валидный. Порядок + важен — `fsync`, потом `replace`, не наоборот; тест проверяет именно + порядок вызовов, а не сам факт вызова `fsync`. + - `load_manifest` читает `manifest.json` обратно в `CollectionManifest`; + отсутствующий или невалидный файл — `ResumeMismatchError`, а не голый + `FileNotFoundError`/pydantic `ValidationError`, чтобы вызывающий код + (резюм) получал понятную доменную ошибку. `CollectionManifest` отклоняет + (через `model_validator`) манифест с двумя `exports` на один и тот же + `WordstatView` — штатная запись (`merge_export`) такое произвести не + может, это защита именно для `load_manifest`, то есть от вручную + отредактированного или битого `manifest.json` на диске, который иначе + молча собьёт `views_to_collect` (возьмёт первую попавшуюся дублирующую + запись). + - `prepare_resume_directory`/`views_to_collect`/`merge_export` — чистая + логика дозаписи существующего run-каталога; подробности и обоснование + главного риска порчи данных (смешение двух разных фраз в одном + каталоге) — в блоке про `collector.py` выше. - **`csv_io.py`** — парсинг только что скачанных CSV. Кодировка автоопределяется перебором (`utf-8-sig`, `utf-8`, `cp1251` — Wordstat @@ -232,8 +315,9 @@ ruff check . - **`errors.py`** — плоская иерархия от `WordstatError`; каждый тип ошибки соответствует конкретной причине сбоя (нет авторизации, изменился UI, - не скачалось вовремя, не распарсился CSV). `cli.py` ловит - `WordstatError | ValueError` и оборачивает в `click.ClickException`. + не скачалось вовремя, не распарсился CSV, `--resume-dir` не подходит — + `ResumeMismatchError`). `cli.py` ловит `WordstatError | ValueError` и + оборачивает в `click.ClickException`. ## Testing notes diff --git a/src/wordstat/cli.py b/src/wordstat/cli.py index b607382..6741b59 100644 --- a/src/wordstat/cli.py +++ b/src/wordstat/cli.py @@ -8,6 +8,7 @@ from wordstat.collector import WordstatCollector from wordstat.config import load_config from wordstat.errors import WordstatError +from wordstat.storage import prepare_resume_directory _CONFIG = load_config() _DEFAULT_CDP_URL = _CONFIG.get("cdp_url", "http://127.0.0.1:9222") @@ -74,6 +75,17 @@ def _read_phrases_file(path: Path) -> str: default=False, help="Keep each downloaded CSV as .csv instead of discarding it after conversion.", ) +@click.option( + "--resume-dir", + "resume_dir", + type=click.Path(path_type=Path, exists=True, file_okay=False), + default=None, + help=( + "Append missing views to an existing run directory instead of starting a new one. " + "Requires exactly one PHRASE, matching the phrase/region recorded in that " + "directory's manifest.json." + ), +) @click.pass_context def collect( ctx: click.Context, @@ -84,6 +96,7 @@ def collect( cdp_url: str, timeout_seconds: float, keep_raw: bool, + resume_dir: Path | None, ) -> None: """Collect all MVP Wordstat reports for one or more PHRASE as Parquet datasets. @@ -95,6 +108,20 @@ def collect( phrases = resolve_phrases(phrase, phrases_file) if not phrases: raise click.ClickException("At least one search phrase is required") + if resume_dir is not None: + if len(phrases) != 1: + raise click.ClickException("--resume-dir requires exactly one phrase") + # Pre-flight check: prepare_resume_directory is pure filesystem logic + # (see storage.py), so a bad --resume-dir (wrong phrase/region, no + # manifest.json, not a directory) can be rejected here before Chrome + # is even touched, instead of surfacing later as a batch failure. + # collect_many/​_collect_one still re-validate the same way right + # before use — this call is a fail-fast convenience, not the only + # guard against a mismatched directory. + try: + prepare_resume_directory(resume_dir, phrases[0], region) + except WordstatError as error: + raise click.ClickException(str(error)) from error collector = WordstatCollector( cdp_url=cdp_url, @@ -103,7 +130,7 @@ def collect( keep_raw=keep_raw, ) try: - batch = asyncio.run(collector.collect_many(phrases, region=region)) + batch = asyncio.run(collector.collect_many(phrases, region=region, resume_directory=resume_dir)) # Only domain errors become friendly messages; an unexpected ValueError # from a dependency should keep its traceback instead of being reworded. except WordstatError as error: diff --git a/src/wordstat/collector.py b/src/wordstat/collector.py index 838c628..543a0d0 100644 --- a/src/wordstat/collector.py +++ b/src/wordstat/collector.py @@ -27,7 +27,14 @@ PhraseFailure, WordstatView, ) -from wordstat.storage import create_run_directory, finalize_raw, write_manifest +from wordstat.storage import ( + create_run_directory, + finalize_raw, + merge_export, + prepare_resume_directory, + views_to_collect, + write_manifest, +) WORDSTAT_URL = "https://wordstat.yandex.ru/" QUERY_SELECTOR = 'input[placeholder="Введите слово или словосочетание"]' @@ -73,10 +80,12 @@ def __init__( self.timeout_seconds = timeout_seconds self.keep_raw = keep_raw - async def collect(self, phrase: str, region: str = "Россия") -> CollectionResult: + async def collect( + self, phrase: str, region: str = "Россия", resume_directory: Path | None = None + ) -> CollectionResult: """Collect popular, related, dynamics and regional reports for one phrase.""" - batch = await self.collect_many([phrase], region=region) + batch = await self.collect_many([phrase], region=region, resume_directory=resume_directory) if batch.failures: raise batch.failures[0].error if not batch.results: @@ -87,7 +96,12 @@ async def collect(self, phrase: str, region: str = "Россия") -> Collection raise InvalidRequestError("Collecting the phrase produced neither a result nor a failure") return batch.results[0] - async def collect_many(self, phrases: list[str], region: str = "Россия") -> BatchCollectionResult: + async def collect_many( + self, + phrases: list[str], + region: str = "Россия", + resume_directory: Path | None = None, + ) -> BatchCollectionResult: """Collect reports for several phrases inside a single browser session. A failure on one phrase is recorded and does not stop the remaining @@ -95,6 +109,12 @@ async def collect_many(self, phrases: list[str], region: str = "Россия") - collected because of one bad phrase. A lost authentication, however, means the session itself is no longer usable, so it aborts the batch instead of repeating the same failure for every remaining phrase. + + ``resume_directory`` requests appending missing views to an existing + run directory instead of creating a new one, and is only valid for a + single phrase: one run directory holds one phrase's manifest, so + resuming while collecting several phrases could not be applied to + all of them without splicing unrelated exports into one manifest. """ region = region.strip() @@ -106,6 +126,8 @@ async def collect_many(self, phrases: list[str], region: str = "Россия") - cleaned_phrases = [phrase.strip() for phrase in phrases] if not all(cleaned_phrases): raise InvalidRequestError("The search phrase must not be empty") + if resume_directory is not None and len(cleaned_phrases) != 1: + raise InvalidRequestError("--resume-dir requires exactly one phrase") results: list[CollectionResult] = [] failures: list[PhraseFailure] = [] @@ -172,6 +194,7 @@ def _mark_region_ready() -> None: region, set_region=not region_ready, on_region_applied=_mark_region_ready, + resume_directory=resume_directory, ) results.append(result) except AuthenticationRequiredError as error: @@ -208,6 +231,7 @@ async def _collect_one( region: str, set_region: bool = True, on_region_applied: Callable[[], None] | None = None, + resume_directory: Path | None = None, ) -> CollectionResult: # Checked once before the batch starts (collect_many), but a session # can lose authentication mid-batch (e.g. Yandex logs it out); check @@ -216,7 +240,27 @@ async def _collect_one( # silently stopped working. await self._assert_authenticated(page) - run_directory = create_run_directory(self.output_root, phrase) + if resume_directory is not None: + # prepare_resume_directory is the hard reject against the main + # data-corruption risk here: it never reuses a directory whose + # stored phrase/region don't match this request. create_run_directory's + # own never-overwrite guarantee is untouched — resuming is this + # separate, explicit path, not a fallback baked into it. + run_directory = resume_directory + manifest = prepare_resume_directory(run_directory, phrase, region) + manifest_path = run_directory / "manifest.json" + pending_views = views_to_collect(run_directory, manifest) + if not pending_views: + return CollectionResult( + run_directory=run_directory, + manifest_path=manifest_path, + manifest=manifest, + ) + else: + run_directory = create_run_directory(self.output_root, phrase) + manifest = None + manifest_path = run_directory / "manifest.json" + pending_views = list(WordstatView) # In a batch, the previous phrase's table can still be sitting in the # DOM when _set_phrase's own waits (new `words=` in the URL, download # button present) are satisfied — those don't check the table itself. @@ -252,8 +296,44 @@ async def _collect_one( required=False, ) - exports = [] - for view, selector in VIEW_SELECTORS.items(): + if manifest is None: + # Written once up front (exports=[], so missing_views/status are + # derived as every view missing / "incomplete") so a crash before + # the first view even finishes still leaves a manifest.json on + # disk describing an empty-but-started run, instead of nothing. + # updated_at is set here too (equal to created_at for this first + # write): leaving it null would make null ambiguous between "a + # successful write already happened" and "this manifest predates + # the updated_at field" — see CollectionManifest's docstring, + # which promises null means only the latter. + start = datetime.now(UTC) + manifest = CollectionManifest( + phrase=phrase, + region=region, + created_at=start, + updated_at=start, + source_url=await page.get_url(), + exports=[], + ) + write_manifest(manifest_path, manifest) + else: + # Resuming: source_url must not silently keep pointing at + # whatever tab/phrase the *previous* session ended on — it is + # only accurate as of the last successful write (see + # CollectionManifest's docstring), and this write is that. + # created_at is deliberately left untouched: it marks when this + # run started, not when every view was collected, and a resumed + # run's views can legitimately span more than one session. + # Doing this after _set_phrase (not before) is required: the + # page is still on the previous phrase/tab until _set_phrase + # returns. + manifest = manifest.model_copy( + update={"source_url": await page.get_url(), "updated_at": datetime.now(UTC)} + ) + write_manifest(manifest_path, manifest) + + for view in pending_views: + selector = VIEW_SELECTORS[view] await self._select_view(page, selector, view) source = await self._download_current_view(page, session, downloads_path) try: @@ -273,25 +353,21 @@ async def _collect_one( if source.exists(): source.replace(run_directory / source.name) raise - exports.append( - ExportSummary( - view=view, - file=data_path.name, - raw_file=raw_path.name if raw_path else None, - row_count=len(dataset.rows), - dtypes=dtypes, - ) + export = ExportSummary( + view=view, + file=data_path.name, + raw_file=raw_path.name if raw_path else None, + row_count=len(dataset.rows), + dtypes=dtypes, ) + # Rewritten after every view, not just once at the end: an + # interruption between views must leave a manifest that honestly + # reflects the views collected so far (see write_manifest and + # CollectionManifest.status), not a stale one still claiming + # every view is missing. + manifest = merge_export(manifest, export) + write_manifest(manifest_path, manifest) - manifest = CollectionManifest( - phrase=phrase, - region=region, - created_at=datetime.now(UTC), - source_url=await page.get_url(), - exports=exports, - ) - manifest_path = run_directory / "manifest.json" - write_manifest(manifest_path, manifest) return CollectionResult( run_directory=run_directory, manifest_path=manifest_path, diff --git a/src/wordstat/errors.py b/src/wordstat/errors.py index eb74f48..0a485c8 100644 --- a/src/wordstat/errors.py +++ b/src/wordstat/errors.py @@ -27,3 +27,7 @@ class DownloadTimeoutError(WordstatError): class CsvFormatError(WordstatError): """A downloaded file cannot be decoded as a headed CSV report.""" + + +class ResumeMismatchError(WordstatError): + """--resume-dir does not belong to the requested phrase/region, or is unusable.""" diff --git a/src/wordstat/models.py b/src/wordstat/models.py index 7394011..92d0102 100644 --- a/src/wordstat/models.py +++ b/src/wordstat/models.py @@ -4,7 +4,7 @@ from enum import StrEnum from pathlib import Path -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, computed_field, model_validator class WordstatView(StrEnum): @@ -46,14 +46,82 @@ class ExportSummary(BaseModel): class CollectionManifest(BaseModel): - """Reproducibility metadata written beside the converted Parquet files.""" + """Reproducibility metadata written beside the converted Parquet files. + + Written incrementally, once per successfully collected view (see + ``collector._collect_one``), so a run interrupted partway through leaves + a manifest on disk that honestly describes what it has so far rather than + none at all. + + ``exports`` is the single source of truth for completeness: both + ``missing_views`` (every :class:`WordstatView` not yet present in + ``exports``, in enum declaration order) and ``status`` are *derived* from + it via ``computed_field`` rather than stored fields, so there is no way + to construct a manifest where they disagree with ``exports`` — the bug + this feature exists to avoid (a caller building + ``CollectionManifest(exports=[])`` without separately remembering to set + ``missing_views`` would otherwise silently get a manifest that lies about + being complete). Both are still plain JSON fields in ``manifest.json`` on + disk (pydantic includes computed fields in ``model_dump_json`` by + default), which is what makes "incomplete" visible to a reader of the + file itself, not only on the in-memory object. + + Three timestamp/URL fields describe when and where the data came from, + and their semantics differ deliberately once ``--resume-dir`` is in + play: + + - ``created_at`` is set once, when the run directory is first created, + and never touched again by a resume. It answers "when did this run + start", not "when was every view collected" — a resumed run's views + can legitimately come from different moments in time. + - ``updated_at`` is the timestamp of the most recent successful write to + this manifest (the initial write or any later resume), so a reader + can tell a fresh run from one stitched together over several + sessions. ``None`` means the manifest predates this field (written by + an older version of this tool) — it is optional so + :func:`~wordstat.storage.load_manifest` can still read a + manifest.json that lacks it, rather than rejecting an otherwise + resumable directory with a validation error. + - ``source_url`` reflects the page URL at the time of the *last* + successful write, not the first: a resume updates it to the current + run's URL (see ``collector._collect_one``), so it never silently + keeps pointing at a stale phrase/tab from a previous session. + + A ``model_validator`` rejects a manifest whose ``exports`` contain more + than one entry for the same :class:`WordstatView`. The normal write path + (``storage.merge_export``) can never produce that — it keeps exports in + a dict keyed by view — but this guards + :func:`~wordstat.storage.load_manifest` against a manifest.json that was + hand-edited or corrupted on disk before a resume reads it back: without + this, ``views_to_collect`` would silently pick whichever duplicate entry + appears first. + """ phrase: str region: str created_at: datetime + updated_at: datetime | None = None source_url: str exports: list[ExportSummary] + @computed_field # type: ignore[prop-decorator] + @property + def missing_views(self) -> list[WordstatView]: + present = {export.view for export in self.exports} + return [view for view in WordstatView if view not in present] + + @computed_field # type: ignore[prop-decorator] + @property + def status(self) -> str: + return "incomplete" if self.missing_views else "complete" + + @model_validator(mode="after") + def _exports_have_unique_views(self) -> "CollectionManifest": + views = [export.view for export in self.exports] + if len(views) != len(set(views)): + raise ValueError("Manifest contains duplicate view exports") + return self + class CollectionResult(BaseModel): """In-memory result returned by :class:`WordstatCollector`. diff --git a/src/wordstat/storage.py b/src/wordstat/storage.py index 86ebea3..4aa0092 100644 --- a/src/wordstat/storage.py +++ b/src/wordstat/storage.py @@ -1,10 +1,13 @@ """Run-directory and manifest handling.""" +import os import re +import tempfile from datetime import UTC, datetime from pathlib import Path -from wordstat.models import CollectionManifest, WordstatView +from wordstat.errors import ResumeMismatchError +from wordstat.models import CollectionManifest, ExportSummary, WordstatView def create_run_directory(output_root: Path, phrase: str, now: datetime | None = None) -> Path: @@ -45,8 +48,137 @@ def finalize_raw(source: Path, run_directory: Path, view: WordstatView, keep_raw def write_manifest(path: Path, manifest: CollectionManifest) -> None: - """Write reproducibility metadata in stable UTF-8 JSON.""" + """Write reproducibility metadata in stable UTF-8 JSON, atomically. - # pydantic emits UTF-8 without escaping non-ASCII, so the Cyrillic stays - # readable without a round-trip through the json module. - path.write_text(manifest.model_dump_json(indent=2) + "\n", encoding="utf-8") + The manifest is written once per phrase after every successfully + collected view (see ``collector._collect_one``), not only at the very + end — a run interrupted mid-batch should leave a manifest that honestly + describes what it managed to collect, marked ``status: "incomplete"`` + (see :class:`~wordstat.models.CollectionManifest`), rather than nothing. + + That means this function can now overwrite an *existing*, previously + valid manifest.json — plain ``Path.write_text`` truncates the file before + writing the new content, so a crash mid-write (process killed, disk full) + would leave a truncated or empty file where a good manifest used to be. + Write to a temporary file in the same directory first, then atomically + replace the target with ``os.replace`` (same filesystem, so it can't fail + partway through) — the target is always either the old content or the + new content, never a half-written mix. ``dir=path.parent`` is required, + not cosmetic: a bare ``tempfile`` call defaults to the system temp + directory, and ``os.replace`` across filesystems raises ``OSError`` + instead of renaming atomically. + + ``os.replace`` only guarantees the *ordering* of the rename, not that the + temporary file's bytes have actually reached disk — on a power loss or + kernel panic the rename can land while the data behind it is still only + in a page cache buffer, leaving an empty or garbage manifest.json where a + valid one used to be. ``flush()`` (drain Python's own buffer) followed by + ``os.fsync()`` (ask the OS to commit it to disk) before the ``replace`` + closes that gap; the whole point of this file is surviving a crash + mid-run, so this one extra syscall is worth it. + """ + + path.parent.mkdir(parents=True, exist_ok=True) + fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=".manifest-", suffix=".json.tmp") + tmp_path = Path(tmp_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + # pydantic emits UTF-8 without escaping non-ASCII, so the + # Cyrillic stays readable without a round-trip through json. + handle.write(manifest.model_dump_json(indent=2) + "\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, path) + finally: + # os.replace already moved the temp file away on the success path, + # so this only fires (and only then removes anything) when writing + # or replace() itself failed partway through. + tmp_path.unlink(missing_ok=True) + + +def load_manifest(path: Path) -> CollectionManifest: + """Read back a previously written manifest.json. + + Raises ``ResumeMismatchError`` if the file is missing or is not valid + manifest JSON, so a caller resuming a run gets a clear domain error + instead of a bare FileNotFoundError/ValidationError. + """ + + if not path.is_file(): + raise ResumeMismatchError(f"No manifest.json found at {path}") + try: + return CollectionManifest.model_validate_json(path.read_text(encoding="utf-8")) + except ValueError as error: + raise ResumeMismatchError(f"{path} is not a valid Wordstat manifest: {error}") from error + + +def prepare_resume_directory(run_directory: Path, phrase: str, region: str) -> CollectionManifest: + """Validate an existing run directory for resuming, and return its manifest. + + A resume directory must already hold a manifest.json for the *same* + request (phrase and region, compared post-strip since both are + normalized before being stored — see collector.collect_many) — otherwise + a typo'd --resume-dir would silently splice a second phrase's exports + into the first phrase's manifest and run directory. This is the primary + data-corruption risk called out for this feature, so the check is a hard + reject (``ResumeMismatchError``), never a best-effort guess. + + A view counts as already collected only when *both* its manifest entry + and its `.parquet` file are present on disk: if the parquet was + deleted by hand after a successful write, the manifest alone must not be + trusted, or resume would report "complete" over a missing file. The + reverse case — a parquet on disk with no manifest entry (a crash between + writing the parquet and rewriting the manifest) — is handled naturally: + the caller re-collects and overwrites that view's file, which is + harmless since the file wasn't accounted for anywhere yet. + """ + + if not run_directory.is_dir(): + raise ResumeMismatchError(f"--resume-dir {run_directory} is not a directory") + + manifest = load_manifest(run_directory / "manifest.json") + if manifest.phrase.strip() != phrase.strip() or manifest.region.strip() != region.strip(): + raise ResumeMismatchError( + f"--resume-dir {run_directory} was collected for phrase={manifest.phrase!r} " + f"region={manifest.region!r}, which does not match the requested " + f"phrase={phrase!r} region={region!r}" + ) + return manifest + + +def views_to_collect(run_directory: Path, manifest: CollectionManifest) -> list[WordstatView]: + """Return the views still missing from a resumed run, in enum order.""" + + done = { + export.view + for export in manifest.exports + if (run_directory / export.file).is_file() + } + return [view for view in WordstatView if view not in done] + + +def merge_export( + manifest: CollectionManifest, export: ExportSummary, now: datetime | None = None +) -> CollectionManifest: + """Return a copy of ``manifest`` with one more view recorded. + + Keeps ``exports`` ordered by :class:`WordstatView` declaration order + (not append order) so a resumed run's manifest looks the same as one + collected in a single pass. ``missing_views``/``status`` are computed + fields derived straight from ``exports`` (see + :class:`~wordstat.models.CollectionManifest`), so updating only + ``exports`` here is enough to keep them correct — there is nothing else + to recompute for those two. + + ``updated_at`` is bumped to ``now`` (defaulting to the current UTC time, + like :func:`create_run_directory` — a caller can pass a fixed value for + deterministic tests) on every call, since every call records a real + write to the manifest. This is a pure function, so it never reads the + clock unless the caller lets it: no hidden ``datetime.now()`` unless + ``now`` is left unset. + """ + + by_view = {item.view: item for item in manifest.exports} + by_view[export.view] = export + exports = [by_view[view] for view in WordstatView if view in by_view] + return manifest.model_copy(update={"exports": exports, "updated_at": now or datetime.now(UTC)}) diff --git a/tests/test_cli.py b/tests/test_cli.py index a33baa8..2e63669 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -66,7 +66,7 @@ def _fake_manifest(phrase: str) -> CollectionManifest: def test_batch_failure_on_one_phrase_does_not_stop_the_rest(monkeypatch, tmp_path: Path): - async def fake_collect_many(self, phrases, region="Россия"): + async def fake_collect_many(self, phrases, region="Россия", resume_directory=None): assert phrases == ["чай", "кофе"] run_directory = tmp_path / "чай" manifest_path = run_directory / "manifest.json" @@ -93,7 +93,7 @@ async def fake_collect_many(self, phrases, region="Россия"): def test_batch_full_success_exits_zero(monkeypatch, tmp_path: Path): - async def fake_collect_many(self, phrases, region="Россия"): + async def fake_collect_many(self, phrases, region="Россия", resume_directory=None): run_directory = tmp_path / "чай" manifest_path = run_directory / "manifest.json" return BatchCollectionResult( @@ -117,7 +117,7 @@ async def fake_collect_many(self, phrases, region="Россия"): def test_batch_aborted_early_reports_untried_phrases_distinctly(monkeypatch): - async def fake_collect_many(self, phrases, region="Россия"): + async def fake_collect_many(self, phrases, region="Россия", resume_directory=None): # Only the first of 3 phrases was attempted (and failed) before the # batch aborted; the CLI must not read this as "2 more failed". return BatchCollectionResult( @@ -154,3 +154,47 @@ def test_phrases_file_with_undecodable_bytes_is_reported_without_a_traceback(tmp assert result.exit_code != 0 assert "Cannot decode --phrases-file" in result.output assert "Traceback" not in result.output + + +def test_resume_dir_with_more_than_one_phrase_is_rejected_before_the_browser_starts(tmp_path: Path): + resume_dir = tmp_path / "some-run" + resume_dir.mkdir() + + result = CliRunner().invoke(main, ["collect", "чай", "кофе", "--resume-dir", str(resume_dir)]) + + assert result.exit_code != 0 + assert "--resume-dir requires exactly one phrase" in result.output + assert "Traceback" not in result.output + + +def test_resume_dir_that_does_not_exist_is_reported_by_click(tmp_path: Path): + missing = tmp_path / "does-not-exist" + + result = CliRunner().invoke(main, ["collect", "чай", "--resume-dir", str(missing)]) + + assert result.exit_code != 0 + assert "Traceback" not in result.output + + +def test_resume_dir_mismatched_phrase_is_rejected_before_the_browser_starts(tmp_path: Path): + from wordstat.models import CollectionManifest as _Manifest + from wordstat.storage import write_manifest + + resume_dir = tmp_path / "some-run" + resume_dir.mkdir() + write_manifest( + resume_dir / "manifest.json", + _Manifest( + phrase="чай", + region="Россия", + created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), + source_url="https://wordstat.yandex.ru/?words=чай", + exports=[], + ), + ) + + result = CliRunner().invoke(main, ["collect", "кофе", "--resume-dir", str(resume_dir)]) + + assert result.exit_code != 0 + assert "does not match" in result.output + assert "Traceback" not in result.output diff --git a/tests/test_collector_batch.py b/tests/test_collector_batch.py index c0f4a93..e7f8f7e 100644 --- a/tests/test_collector_batch.py +++ b/tests/test_collector_batch.py @@ -13,8 +13,9 @@ import wordstat.collector as collector_module from wordstat.collector import WordstatCollector -from wordstat.errors import AuthenticationRequiredError, InvalidRequestError, PhraseEntryError -from wordstat.models import CollectionManifest, CollectionResult +from wordstat.errors import AuthenticationRequiredError, InvalidRequestError, PhraseEntryError, ResumeMismatchError +from wordstat.models import CollectionManifest, CollectionResult, WordstatView +from wordstat.storage import load_manifest def _fake_result(tmp_path, phrase) -> CollectionResult: @@ -33,11 +34,14 @@ def _fake_result(tmp_path, phrase) -> CollectionResult: class _FakePage: + def __init__(self, url: str = "https://wordstat.yandex.ru/?words=test"): + self.url = url + async def goto(self, url): pass async def get_url(self): - return "https://wordstat.yandex.ru/?words=test" + return self.url class _FakeSession: @@ -73,7 +77,15 @@ def _patch_common(monkeypatch): def _patch_collect_one(monkeypatch, behavior): async def fake_collect_one( - self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, ): return await behavior(phrase) @@ -142,7 +154,17 @@ def test_collect_many_aborts_on_lost_authentication_without_trying_the_rest(monk _patch_common(monkeypatch) attempted = [] - async def collect_one(self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None): + async def collect_one( + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, + ): attempted.append(phrase) if phrase == "второй": raise AuthenticationRequiredError("session lost") @@ -186,7 +208,17 @@ def _patch_collect_one_passthrough(monkeypatch, tmp_path): """Let _collect_one call the real _assert_authenticated, then short-circuit the rest of the browser interaction.""" - async def passthrough(self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None): + async def passthrough( + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, + ): await self._assert_authenticated(page) return _fake_result(tmp_path, phrase) @@ -205,7 +237,15 @@ def test_collect_many_only_sets_region_for_the_first_phrase(monkeypatch, tmp_pat set_region_calls = [] async def recording_collect_one( - self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, ): set_region_calls.append(set_region) if set_region and on_region_applied is not None: @@ -233,7 +273,15 @@ def test_collect_many_retries_region_after_the_first_phrase_fails(monkeypatch, t set_region_calls = [] async def recording_collect_one( - self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, ): set_region_calls.append(set_region) if phrase == "первый": @@ -267,7 +315,15 @@ def test_collect_many_does_not_retry_region_if_it_was_applied_before_a_later_fai set_region_calls = [] async def collect_one( - self, page, session, downloads_path, phrase, region, set_region=True, on_region_applied=None + self, + page, + session, + downloads_path, + phrase, + region, + set_region=True, + on_region_applied=None, + resume_directory=None, ): set_region_calls.append(set_region) if set_region: @@ -455,3 +511,311 @@ async def run(): # Must raise the original RuntimeError, not a FileNotFoundError from the # rescue trying to move an already-moved file. asyncio.run(run()) + + +# --- incremental manifest / resume, exercised through the real _collect_one --- + + +def test_collect_one_writes_the_manifest_after_every_view_not_only_at_the_end(monkeypatch, tmp_path): + """The manifest.json on disk must reflect progress mid-phrase, not just + the final result — otherwise a crash between views leaves nothing to + resume from.""" + + _patch_common(monkeypatch) + + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + call_count = {"n": 0} + + async def fake_download(self, page, session, dl_path): + call_count["n"] += 1 + source = dl_path / f"export-{call_count['n']}.csv" + _write_view_csv(source, "тест") + return source + + seen_statuses = [] + real_write_manifest = collector_module.write_manifest + + def recording_write_manifest(path, manifest): + real_write_manifest(path, manifest) + seen_statuses.append((manifest.status, len(manifest.exports))) + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + monkeypatch.setattr(collector_module, "write_manifest", recording_write_manifest) + + collector = WordstatCollector("cdp", tmp_path) + + async def run(): + page = _FakePage() + session = _FakeSession() + return await collector._collect_one(page, session, downloads_path, "тест", "Россия") + + result = asyncio.run(run()) + + # One write before the loop (empty, incomplete) + one per view (4). + assert seen_statuses[0] == ("incomplete", 0) + assert seen_statuses[-1] == ("complete", 4) + assert len(seen_statuses) == 5 + assert result.manifest.status == "complete" + on_disk = load_manifest(result.manifest_path) + assert on_disk.status == "complete" + assert len(on_disk.exports) == 4 + # Even the very first write (before any view exists) must set + # updated_at, not leave it null — null is reserved for manifests written + # by a version of this tool that predates the field, and a fresh run + # must never look like one of those. + assert on_disk.updated_at is not None + assert on_disk.updated_at >= on_disk.created_at + + +def test_collect_one_resume_directory_only_collects_missing_views(monkeypatch, tmp_path): + """Resuming an existing run directory must add only the missing views + and must not re-download or overwrite views already recorded with their + file present on disk.""" + + _patch_common(monkeypatch) + + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + # First pass: run for real, but make it fail right after the first view + # so the run directory is left genuinely incomplete. + async def fake_select_view(self, page, selector, view): + pass + + call_count = {"n": 0} + + async def failing_after_first_download(self, page, session, dl_path): + call_count["n"] += 1 + if call_count["n"] > 1: + raise RuntimeError("simulated interruption") + source = dl_path / "export-1.csv" + _write_view_csv(source, "тест") + return source + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", failing_after_first_download) + + collector = WordstatCollector("cdp", tmp_path) + + async def run_first(): + page = _FakePage() + session = _FakeSession() + with pytest.raises(RuntimeError, match="simulated interruption"): + await collector._collect_one(page, session, downloads_path, "тест", "Россия") + + asyncio.run(run_first()) + + run_directories = [p for p in tmp_path.glob("runs/*") if p.is_dir()] + assert len(run_directories) == 1 + run_directory = run_directories[0] + partial_manifest = load_manifest(run_directory / "manifest.json") + assert partial_manifest.status == "incomplete" + first_view_export = next(e for e in partial_manifest.exports if e.view == WordstatView.TOP_POPULAR) + first_view_mtime = (run_directory / first_view_export.file).stat().st_mtime + + # Second pass: resume, and this time let every remaining view succeed. + call_count["n"] = 0 + + async def succeeding_download(self, page, session, dl_path): + call_count["n"] += 1 + source = dl_path / f"resume-{call_count['n']}.csv" + _write_view_csv(source, "тест") + return source + + monkeypatch.setattr(WordstatCollector, "_download_current_view", succeeding_download) + + async def run_resume(): + page = _FakePage() + session = _FakeSession() + return await collector._collect_one( + page, session, downloads_path, "тест", "Россия", resume_directory=run_directory + ) + + result = asyncio.run(run_resume()) + + assert result.run_directory == run_directory + assert call_count["n"] == 3 # only the 3 previously-missing views were downloaded + final_manifest = load_manifest(run_directory / "manifest.json") + assert final_manifest.status == "complete" + assert len(final_manifest.exports) == 4 + # The already-collected view's file was not touched by the resume. + still_there = next(e for e in final_manifest.exports if e.view == WordstatView.TOP_POPULAR) + assert still_there == first_view_export + assert (run_directory / still_there.file).stat().st_mtime == first_view_mtime + + +def test_collect_one_resume_directory_rejects_a_different_phrase(monkeypatch, tmp_path): + _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_view_csv(source, "чай") + return source + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path) + + async def run(): + page = _FakePage() + session = _FakeSession() + return await collector._collect_one(page, session, downloads_path, "чай", "Россия") + + result = asyncio.run(run()) + + async def run_resume_with_wrong_phrase(): + page = _FakePage() + session = _FakeSession() + with pytest.raises(ResumeMismatchError, match="does not match"): + await collector._collect_one( + page, + session, + downloads_path, + "кофе", + "Россия", + resume_directory=result.run_directory, + ) + + asyncio.run(run_resume_with_wrong_phrase()) + + +def test_collect_many_rejects_resume_directory_with_more_than_one_phrase(tmp_path): + collector = WordstatCollector("cdp", tmp_path) + resume_dir = tmp_path / "some-run" + resume_dir.mkdir() + + with pytest.raises(InvalidRequestError, match="--resume-dir requires exactly one phrase"): + asyncio.run(collector.collect_many(["чай", "кофе"], resume_directory=resume_dir)) + + +# --- resume updates source_url/updated_at, leaves created_at alone (bugfix #3) --- + + +def test_collect_one_resume_updates_source_url_and_updated_at_but_not_created_at(monkeypatch, tmp_path): + """The bug no reviewer caught: resuming used to leave created_at and + source_url exactly as they were after the first, interrupted pass, even + though the remaining views are collected later, from a different page + state. created_at must still describe when the run *started*; source_url + must reflect the page at the time of this write, not the first one.""" + + _patch_common(monkeypatch) + + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + + async def fake_select_view(self, page, selector, view): + pass + + async def failing_download(self, page, session, dl_path): + raise RuntimeError("simulated interruption before any view finished") + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", failing_download) + + collector = WordstatCollector("cdp", tmp_path) + + async def run_first(): + page = _FakePage(url="https://wordstat.yandex.ru/?words=тест®ion=Россия") + session = _FakeSession() + with pytest.raises(RuntimeError, match="simulated interruption"): + await collector._collect_one(page, session, downloads_path, "тест", "Россия") + + asyncio.run(run_first()) + + run_directory = next(p for p in tmp_path.glob("runs/*") if p.is_dir()) + first_manifest = load_manifest(run_directory / "manifest.json") + original_created_at = first_manifest.created_at + original_source_url = first_manifest.source_url + original_updated_at = first_manifest.updated_at + + async def succeeding_download(self, page, session, dl_path): + source = dl_path / "export.csv" + _write_view_csv(source, "тест") + return source + + monkeypatch.setattr(WordstatCollector, "_download_current_view", succeeding_download) + + resumed_url = "https://wordstat.yandex.ru/?words=тест®ion=Россия&tab=table" + + async def run_resume(): + page = _FakePage(url=resumed_url) + session = _FakeSession() + return await collector._collect_one( + page, session, downloads_path, "тест", "Россия", resume_directory=run_directory + ) + + result = asyncio.run(run_resume()) + + assert result.manifest.created_at == original_created_at # start time is untouched + assert result.manifest.source_url == resumed_url # reflects this session's page, not the old one + assert result.manifest.source_url != original_source_url + # The very first write already sets updated_at (equal to created_at at + # that point) — see _collect_one — so null is unambiguous elsewhere as + # "predates this field". A resume must bump it forward, never leave it + # equal to (let alone before) the interrupted first pass's value. + assert original_updated_at is not None + assert result.manifest.updated_at is not None + assert result.manifest.updated_at > original_updated_at + + on_disk = load_manifest(result.manifest_path) + assert on_disk.created_at == original_created_at + assert on_disk.source_url == resumed_url + + +def test_collect_one_resume_of_an_already_complete_run_does_not_touch_the_manifest(monkeypatch, tmp_path): + """A resume that finds nothing missing (the early return in + _collect_one) must not bump updated_at or rewrite source_url — nothing + was actually collected in that call, so claiming an update would be a + lie.""" + + _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_view_csv(source, "тест") + return source + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + + collector = WordstatCollector("cdp", tmp_path) + + async def run_full(): + page = _FakePage(url="https://wordstat.yandex.ru/?words=тест") + session = _FakeSession() + return await collector._collect_one(page, session, downloads_path, "тест", "Россия") + + first_result = asyncio.run(run_full()) + manifest_bytes_before = first_result.manifest_path.read_bytes() + + async def run_resume(): + page = _FakePage(url="https://wordstat.yandex.ru/?words=тест&tab=different") + session = _FakeSession() + return await collector._collect_one( + page, session, downloads_path, "тест", "Россия", resume_directory=first_result.run_directory + ) + + resumed_result = asyncio.run(run_resume()) + + assert resumed_result.manifest.source_url == first_result.manifest.source_url + assert resumed_result.manifest.updated_at == first_result.manifest.updated_at + assert resumed_result.manifest_path.read_bytes() == manifest_bytes_before diff --git a/tests/test_storage.py b/tests/test_storage.py index 39ee787..f85ecb0 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -1,8 +1,22 @@ +import json +import os from datetime import UTC, datetime from pathlib import Path +import pytest + +from wordstat.errors import ResumeMismatchError from wordstat.models import CollectionManifest, ExportSummary, WordstatView -from wordstat.storage import create_run_directory, finalize_raw, slugify, write_manifest +from wordstat.storage import ( + create_run_directory, + finalize_raw, + load_manifest, + merge_export, + prepare_resume_directory, + slugify, + views_to_collect, + write_manifest, +) def test_create_run_directory_is_unique_and_keeps_cyrillic(tmp_path: Path) -> None: @@ -66,3 +80,378 @@ def test_finalize_raw_renames_the_download_when_keeping_it(tmp_path: Path) -> No def test_finalize_raw_tolerates_an_already_missing_download(tmp_path: Path) -> None: assert finalize_raw(tmp_path / "gone.csv", tmp_path, WordstatView.DYNAMICS, keep_raw=False) is None + + +def _manifest( + phrase: str = "ремонт квартир", + region: str = "Москва", + exports: list[ExportSummary] | None = None, +) -> CollectionManifest: + return CollectionManifest( + phrase=phrase, + region=region, + created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), + source_url="https://wordstat.yandex.ru/?words=test", + exports=exports if exports is not None else [], + ) + + +def _export(view: WordstatView, file: str | None = None) -> ExportSummary: + return ExportSummary( + view=view, + file=file or f"{view.value}.parquet", + raw_file=None, + row_count=1, + dtypes={"Запрос": "string"}, + ) + + +# --- status / missing_views visible on disk ------------------------------- + + +def test_manifest_json_marks_an_incomplete_run_on_disk(tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + manifest = _manifest(exports=[_export(WordstatView.TOP_POPULAR)]) + + write_manifest(path, manifest) + + on_disk = json.loads(path.read_text(encoding="utf-8")) + assert on_disk["status"] == "incomplete" + assert on_disk["missing_views"] == [ + v.value for v in WordstatView if v != WordstatView.TOP_POPULAR + ] + + +def test_manifest_json_marks_a_complete_run_on_disk(tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + manifest = _manifest(exports=[_export(view) for view in WordstatView]) + + write_manifest(path, manifest) + + on_disk = json.loads(path.read_text(encoding="utf-8")) + assert on_disk["status"] == "complete" + assert on_disk["missing_views"] == [] + + +def test_manifest_status_and_missing_views_cannot_disagree_with_exports() -> None: + # Regression guard for the original design flaw: status/missing_views + # used to be separate stored fields that a caller could set + # inconsistently with exports (e.g. exports=[] with no missing_views, + # which would have reported "complete" on zero exports). Both are now + # computed straight from exports, so there is no constructor call that + # can produce a disagreement. + manifest = _manifest(exports=[]) + + assert manifest.status == "incomplete" + assert manifest.missing_views == list(WordstatView) + + +# --- atomic write ------------------------------------------------------ + + +def test_write_manifest_leaves_the_previous_content_intact_if_replace_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "manifest.json" + original = _manifest(phrase="исходная фраза") + write_manifest(path, original) + original_bytes = path.read_bytes() + + def failing_replace(*_args, **_kwargs): + raise OSError("simulated crash between write and replace") + + monkeypatch.setattr("wordstat.storage.os.replace", failing_replace) + + with pytest.raises(OSError, match="simulated crash"): + write_manifest(path, _manifest(phrase="новая фраза")) + + # The target file was never truncated: it's still exactly the old, + # valid manifest, not empty and not a half-written mix. + assert path.read_bytes() == original_bytes + assert CollectionManifest.model_validate_json(path.read_text(encoding="utf-8")).phrase == "исходная фраза" + + +def test_write_manifest_leaves_no_temp_files_behind_on_success(tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + + write_manifest(path, _manifest()) + + assert [p.name for p in tmp_path.iterdir()] == ["manifest.json"] + + +def test_write_manifest_leaves_no_temp_files_behind_on_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + path = tmp_path / "manifest.json" + write_manifest(path, _manifest()) + + def failing_replace(*_args, **_kwargs): + raise OSError("simulated crash") + + monkeypatch.setattr("wordstat.storage.os.replace", failing_replace) + with pytest.raises(OSError): + write_manifest(path, _manifest(phrase="другая фраза")) + + assert [p.name for p in tmp_path.iterdir()] == ["manifest.json"] + + +def test_write_manifest_writes_the_temp_file_next_to_the_target(tmp_path: Path) -> None: + # Regression guard: a bare tempfile.mkstemp()/NamedTemporaryFile() call + # without dir= defaults to the system temp directory, which can sit on a + # different filesystem than --output-dir. os.replace() across + # filesystems raises OSError instead of renaming atomically, so the temp + # file must be created in path.parent. + nested = tmp_path / "runs" / "some-run" + path = nested / "manifest.json" + + write_manifest(path, _manifest()) + + assert path.exists() + assert [p.name for p in nested.iterdir()] == ["manifest.json"] + + +# --- resume: mismatch rejection ----------------------------------------- + + +def test_prepare_resume_directory_rejects_a_directory_without_a_manifest(tmp_path: Path) -> None: + run_directory = tmp_path / "runs" / "empty" + run_directory.mkdir(parents=True) + + with pytest.raises(ResumeMismatchError, match="No manifest.json"): + prepare_resume_directory(run_directory, "ремонт квартир", "Москва") + + +def test_prepare_resume_directory_rejects_a_non_directory(tmp_path: Path) -> None: + not_a_dir = tmp_path / "not-a-dir" + not_a_dir.write_text("oops", encoding="utf-8") + + with pytest.raises(ResumeMismatchError, match="is not a directory"): + prepare_resume_directory(not_a_dir, "ремонт квартир", "Москва") + + +def test_prepare_resume_directory_rejects_corrupt_json(tmp_path: Path) -> None: + run_directory = tmp_path / "runs" / "broken" + run_directory.mkdir(parents=True) + (run_directory / "manifest.json").write_text("{not valid json", encoding="utf-8") + + with pytest.raises(ResumeMismatchError, match="not a valid Wordstat manifest"): + prepare_resume_directory(run_directory, "ремонт квартир", "Москва") + + +def test_prepare_resume_directory_rejects_a_different_phrase(tmp_path: Path) -> None: + run_directory = tmp_path / "runs" / "run" + run_directory.mkdir(parents=True) + write_manifest(run_directory / "manifest.json", _manifest(phrase="ремонт квартир")) + + with pytest.raises(ResumeMismatchError, match="does not match"): + prepare_resume_directory(run_directory, "натяжные потолки", "Москва") + + +def test_prepare_resume_directory_rejects_a_different_region(tmp_path: Path) -> None: + run_directory = tmp_path / "runs" / "run" + run_directory.mkdir(parents=True) + write_manifest(run_directory / "manifest.json", _manifest(region="Москва")) + + with pytest.raises(ResumeMismatchError, match="does not match"): + prepare_resume_directory(run_directory, "ремонт квартир", "Санкт-Петербург") + + +def test_prepare_resume_directory_accepts_whitespace_differences(tmp_path: Path) -> None: + # collect_many strips phrase/region before storing them in the manifest, + # so the comparison here must also compare stripped values, or a + # trailing space on the CLI argument would cause a false rejection. + run_directory = tmp_path / "runs" / "run" + run_directory.mkdir(parents=True) + write_manifest(run_directory / "manifest.json", _manifest(phrase="ремонт квартир", region="Москва")) + + manifest = prepare_resume_directory(run_directory, " ремонт квартир ", " Москва ") + + assert manifest.phrase == "ремонт квартир" + + +# --- resume: views_to_collect / merge_export ------------------------------ + + +def test_views_to_collect_skips_views_recorded_with_their_file_present(tmp_path: Path) -> None: + run_directory = tmp_path + (run_directory / "top_popular.parquet").write_bytes(b"data") + manifest = _manifest(exports=[_export(WordstatView.TOP_POPULAR)]) + + pending = views_to_collect(run_directory, manifest) + + assert pending == [v for v in WordstatView if v != WordstatView.TOP_POPULAR] + + +def test_views_to_collect_does_not_trust_a_manifest_entry_whose_file_was_deleted(tmp_path: Path) -> None: + # The parquet for top_popular was removed by hand after a successful + # write; the manifest entry alone must not be enough to skip it again. + run_directory = tmp_path + manifest = _manifest(exports=[_export(WordstatView.TOP_POPULAR)]) + + pending = views_to_collect(run_directory, manifest) + + assert pending == list(WordstatView) + + +def test_merge_export_adds_a_missing_view_without_touching_existing_ones() -> None: + existing = _export(WordstatView.TOP_POPULAR) + manifest = _manifest(exports=[existing]) + + updated = merge_export(manifest, _export(WordstatView.REGIONS)) + + by_view = {item.view: item for item in updated.exports} + assert by_view[WordstatView.TOP_POPULAR] == existing + assert WordstatView.REGIONS in by_view + assert updated.missing_views == [WordstatView.TOP_RELATED, WordstatView.DYNAMICS] + + +def test_merge_export_orders_exports_by_view_declaration_not_append_order() -> None: + manifest = _manifest(exports=[]) + + manifest = merge_export(manifest, _export(WordstatView.REGIONS)) + manifest = merge_export(manifest, _export(WordstatView.TOP_POPULAR)) + + assert [item.view for item in manifest.exports] == [WordstatView.TOP_POPULAR, WordstatView.REGIONS] + + +def test_merge_export_marks_the_manifest_complete_once_every_view_is_present() -> None: + manifest = _manifest(exports=[]) + + for view in WordstatView: + manifest = merge_export(manifest, _export(view)) + + assert manifest.missing_views == [] + assert manifest.status == "complete" + + +# --- load_manifest --------------------------------------------------------- + + +def test_load_manifest_round_trips_what_write_manifest_wrote(tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + original = _manifest(phrase="ремонт квартир", exports=[_export(WordstatView.TOP_POPULAR)]) + write_manifest(path, original) + + loaded = load_manifest(path) + + assert loaded.phrase == "ремонт квартир" + assert [item.view for item in loaded.exports] == [WordstatView.TOP_POPULAR] + + +def test_load_manifest_rejects_a_missing_file(tmp_path: Path) -> None: + with pytest.raises(ResumeMismatchError, match="No manifest.json"): + load_manifest(tmp_path / "manifest.json") + + +# --- fsync before os.replace (write_manifest atomicity, cheri-pick #1) ----- + + +def test_write_manifest_fsyncs_before_replacing_the_target(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + # A test that only asserts fsync was *called* would also pass if fsync + # ran after os.replace, where it can no longer prevent a torn write. + # Record the actual order both land in. + path = tmp_path / "manifest.json" + order = [] + + real_fsync = os.fsync + real_replace = os.replace + + def recording_fsync(fd): + order.append("fsync") + return real_fsync(fd) + + def recording_replace(src, dst): + order.append("replace") + return real_replace(src, dst) + + monkeypatch.setattr("wordstat.storage.os.fsync", recording_fsync) + monkeypatch.setattr("wordstat.storage.os.replace", recording_replace) + + write_manifest(path, _manifest()) + + assert order == ["fsync", "replace"] + + +# --- unique-view validator (cheri-pick #2) --------------------------------- + + +def test_collection_manifest_rejects_duplicate_view_exports() -> None: + with pytest.raises(ValueError, match="duplicate view exports"): + CollectionManifest( + phrase="ремонт квартир", + region="Москва", + created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), + source_url="https://wordstat.yandex.ru/?words=test", + exports=[_export(WordstatView.DYNAMICS), _export(WordstatView.DYNAMICS)], + ) + + +def test_load_manifest_rejects_a_hand_edited_file_with_duplicate_views(tmp_path: Path) -> None: + # The threat this guards against is a corrupted/hand-edited file on + # disk, read back by a resume — not a constructor call in this + # codebase's own code (merge_export can't produce this). ValidationError + # is a ValueError subclass, so load_manifest's existing `except + # ValueError` must turn it into the same ResumeMismatchError a resuming + # caller already handles. + path = tmp_path / "manifest.json" + write_manifest(path, _manifest(exports=[_export(WordstatView.DYNAMICS)])) + raw = json.loads(path.read_text(encoding="utf-8")) + raw["exports"].append(dict(raw["exports"][0])) # duplicate the one export + path.write_text(json.dumps(raw, ensure_ascii=False, indent=2), encoding="utf-8") + + with pytest.raises(ResumeMismatchError, match="not a valid Wordstat manifest"): + load_manifest(path) + + +def test_merge_export_does_not_reintroduce_duplicate_views() -> None: + # merge_export replaces by view in a dict, so re-merging the same view + # must not somehow produce two entries — confirms the validator and the + # write path agree rather than one silently working around the other. + manifest = _manifest(exports=[_export(WordstatView.DYNAMICS)]) + + updated = merge_export(manifest, _export(WordstatView.DYNAMICS, file="dynamics-v2.parquet")) + + assert [e.view for e in updated.exports] == [WordstatView.DYNAMICS] + assert updated.exports[0].file == "dynamics-v2.parquet" + + +# --- updated_at optionality / round-trip (prerequisite for bugfix #3) ------ + + +def test_updated_at_defaults_to_none() -> None: + assert _manifest().updated_at is None + + +def test_load_manifest_accepts_a_manifest_json_written_before_updated_at_existed(tmp_path: Path) -> None: + # A manifest.json on disk from an older version of this tool has no + # "updated_at" key at all. load_manifest must still validate it (as + # None), not reject an otherwise perfectly resumable directory. + path = tmp_path / "manifest.json" + payload = json.loads(_manifest(exports=[_export(WordstatView.TOP_POPULAR)]).model_dump_json()) + del payload["updated_at"] + path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + + loaded = load_manifest(path) + + assert loaded.updated_at is None + + +# --- merge_export bumps updated_at (prerequisite for bugfix #3) ------------ + + +def test_merge_export_sets_updated_at_to_the_given_now() -> None: + now = datetime(2026, 8, 21, 9, 30, tzinfo=UTC) + manifest = _manifest(exports=[]) + + updated = merge_export(manifest, _export(WordstatView.TOP_POPULAR), now=now) + + assert updated.updated_at == now + + +def test_merge_export_does_not_touch_created_at() -> None: + manifest = _manifest(exports=[]) + original_created_at = manifest.created_at + + updated = merge_export(manifest, _export(WordstatView.TOP_POPULAR), now=datetime(2026, 8, 22, tzinfo=UTC)) + + assert updated.created_at == original_created_at