tests/unit/application/test_annotate.py

# License: Apache-2.0 License
# Created by: Patrick Lumbantobing, VertoX-AI
# Copyright (c) 2026 VertoX-AI. All rights reserved.
#
# This work is licensed under the Apache-2.0 License.
# To view a copy of this license, visit
# https://www.apache.org/licenses/LICENSE-2.0

from __future__ import annotations

import json
import os
from collections.abc import Sequence
from pathlib import Path

import pytest

from speech_text_data_aligners.application.annotate import (
    AnnotationOptions,
    _write_failures,
    annotate_metadata,
)
from speech_text_data_aligners.core.errors import (
    BackendExecutionError,
    BackendOutputError,
    ConfigurationError,
    SpeechTextAlignerError,
)
from speech_text_data_aligners.core.models import (
    AlignmentFailure,
    AlignmentOutcome,
    AlignmentProvenance,
    AlignmentRequest,
    AlignmentResult,
    WordAlignment,
)
from speech_text_data_aligners.core.normalization import lexicalize
from speech_text_data_aligners.data.metadata import read_metadata


class FakeBackend:
    def __init__(
        self,
        *,
        fail_ids: set[str] | None = None,
        raise_error: bool = False,
        bad_count: bool = False,
        wrong_id: bool = False,
        invalid_timestamps: bool = False,
        offset: float = 0.0,
    ) -> None:
        self.fail_ids = fail_ids or set()
        self.raise_error = raise_error
        self.bad_count = bad_count
        self.wrong_id = wrong_id
        self.invalid_timestamps = invalid_timestamps
        self.offset = offset
        self.batches: list[tuple[AlignmentRequest, ...]] = []

    @property
    def name(self) -> str:
        return "fake"

    def _result(self, request: AlignmentRequest) -> AlignmentResult:
        words = lexicalize(request.transcript)
        alignments = tuple(
            WordAlignment(
                word,
                self.offset + index * 0.3,
                (2.0 if self.invalid_timestamps else self.offset + index * 0.3 + 0.2),
            )
            for index, word in enumerate(words)
        )
        return AlignmentResult(
            "wrong" if self.wrong_id else request.request_id,
            alignments,
            AlignmentProvenance("fake", backend_version="1"),
            runtime_s=0.1,
        )

    def align(self, request: AlignmentRequest) -> AlignmentResult:
        return self._result(request)

    def align_batch(
        self, requests: Sequence[AlignmentRequest]
    ) -> tuple[AlignmentOutcome, ...]:
        batch = tuple(requests)
        self.batches.append(batch)
        if self.raise_error:
            raise BackendExecutionError("provider unavailable", retryable=True)
        outcomes = tuple(
            AlignmentOutcome(
                failure=AlignmentFailure(
                    request.request_id, "fake_failure", "test", "failed"
                )
            )
            if request.request_id in self.fail_ids
            else AlignmentOutcome(result=self._result(request))
            for request in batch
        )
        return outcomes[:-1] if self.bad_count else outcomes


def test_annotation_success_preserves_rows_and_writes_optional_receipt(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    before = metadata.read_bytes()
    backend = FakeBackend()
    output = tmp_path / "annotated.csv"
    receipt = tmp_path / "receipt.json"
    run = annotate_metadata(
        metadata_csv=metadata,
        audio_dir=audio_dir,
        output_csv=output,
        backend=backend,
        options=AnnotationOptions(receipt_path=receipt),
    )
    table = read_metadata(output)
    assert metadata.read_bytes() == before
    assert [table.row_dict(i)["en_transcript"] for i in range(2)] == [
        "Hello world.",
        "Can't stop",
    ]
    assert table.row_dict(0)["first_word_text"] == "Hello"
    assert table.row_dict(1)["first_word_text"] == "Can't"
    assert len(backend.batches) == 1
    assert [request.request_id for request in backend.batches[0]] == [
        "000000",
        "000001",
    ]
    assert run.failures == ()
    assert run.receipt.success_count == 2
    assert run.receipt.failure_count == 0
    assert json.loads(receipt.read_text(encoding="utf-8"))["operation"] == "annotate"


def test_continue_mode_blanks_failed_rows_and_writes_diagnostics(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    output = tmp_path / "annotated.csv"
    run = annotate_metadata(
        metadata_csv=metadata,
        audio_dir=audio_dir,
        output_csv=output,
        backend=FakeBackend(fail_ids={"000001"}),
    )
    table = read_metadata(output)
    failed_row = table.row_dict(1)
    assert failed_row["first_word_text"] == ""
    assert failed_row["last_word_end_s"] == ""
    assert run.diagnostics_path == tmp_path / "annotated.diagnostics.json"
    diagnostics = json.loads(run.diagnostics_path.read_text(encoding="utf-8"))
    assert diagnostics[0]["request_id"] == "000001"
    assert run.failures[0].code == "fake_failure"


def test_continue_mode_structures_preflight_and_whole_backend_failures(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    (audio_dir / "a2.wav").unlink()
    backend = FakeBackend(raise_error=True)
    run = annotate_metadata(
        metadata_csv=metadata,
        audio_dir=audio_dir,
        output_csv=tmp_path / "out.csv",
        backend=backend,
    )
    assert [failure.code for failure in run.failures] == [
        "backend_execution_error",
        "audio_error",
    ]
    assert len(backend.batches[0]) == 1


def test_continue_mode_converts_bad_count_wrong_identity_and_invalid_timestamp(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    for index, backend in enumerate(
        (
            FakeBackend(bad_count=True),
            FakeBackend(wrong_id=True),
            FakeBackend(invalid_timestamps=True),
        )
    ):
        run = annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=tmp_path / f"out{index}.csv",
            backend=backend,
        )
        assert len(run.failures) == 2
    assert run.failures[0].code == "alignment_validation_error"


def test_strict_mode_aborts_without_output_on_row_or_backend_failure(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    output = tmp_path / "out.csv"
    with pytest.raises(SpeechTextAlignerError):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=output,
            backend=FakeBackend(fail_ids={"000000"}),
            options=AnnotationOptions(strict=True),
        )
    assert not output.exists()

    (audio_dir / "a1.wav").unlink()
    with pytest.raises(SpeechTextAlignerError):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=output,
            backend=FakeBackend(),
            options=AnnotationOptions(strict=True),
        )
    assert not output.exists()


def test_annotation_handles_no_valid_requests(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    (audio_dir / "a1.wav").unlink()
    (audio_dir / "a2.wav").unlink()
    backend = FakeBackend()
    run = annotate_metadata(
        metadata_csv=metadata,
        audio_dir=audio_dir,
        output_csv=tmp_path / "out.csv",
        backend=backend,
    )
    assert len(run.failures) == 2
    assert backend.batches == []


def test_annotation_preflights_existing_diagnostics_before_output(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    diagnostics = tmp_path / "existing.json"
    diagnostics.write_text("old", encoding="utf-8")
    output = tmp_path / "out.csv"
    backend = FakeBackend(fail_ids={"000000"})
    with pytest.raises(ConfigurationError, match="exists"):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=output,
            backend=backend,
            options=AnnotationOptions(diagnostics_path=diagnostics),
        )
    assert not output.exists()
    assert backend.batches == []


@pytest.mark.parametrize(
    "collision",
    [
        "output-source",
        "diagnostics-source",
        "receipt-source",
        "diagnostics-output",
        "receipt-output",
        "receipt-diagnostics",
    ],
)
def test_annotation_rejects_all_artifact_collisions_before_backend_or_mutation(
    metadata_fixture: tuple[Path, Path], tmp_path: Path, collision: str
) -> None:
    metadata, audio_dir = metadata_fixture
    output = tmp_path / "out.csv"
    diagnostics = tmp_path / "diagnostics.json"
    receipt = tmp_path / "receipt.json"
    if collision == "output-source":
        output = metadata
    elif collision == "diagnostics-source":
        diagnostics = metadata
    elif collision == "receipt-source":
        receipt = metadata
    elif collision == "diagnostics-output":
        diagnostics = output
    elif collision == "receipt-output":
        receipt = output
    else:
        receipt = diagnostics

    for path in {output, diagnostics, receipt} - {metadata}:
        path.write_bytes(f"preserve:{path.name}".encode())
    snapshots = {
        path: path.read_bytes() for path in {metadata, output, diagnostics, receipt}
    }
    backend = FakeBackend()
    with pytest.raises(ConfigurationError, match="paths alias"):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=output,
            backend=backend,
            options=AnnotationOptions(
                overwrite=True,
                diagnostics_path=diagnostics,
                receipt_path=receipt,
            ),
        )
    assert backend.batches == []
    assert {path: path.read_bytes() for path in snapshots} == snapshots


def test_annotation_rejects_normalized_relative_alias_before_backend(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    relative_alias = Path(
        os.path.relpath(metadata.parent / "not-created" / ".." / metadata.name)
    )
    backend = FakeBackend()
    before = metadata.read_bytes()
    with pytest.raises(ConfigurationError, match="paths alias"):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=relative_alias,
            backend=backend,
            options=AnnotationOptions(overwrite=True),
        )
    assert backend.batches == []
    assert metadata.read_bytes() == before


def test_annotation_rejects_existing_symlink_alias_before_backend(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    alias = tmp_path / "metadata-alias.csv"
    try:
        alias.symlink_to(metadata)
    except OSError as error:
        pytest.skip(f"symlinks are unavailable: {error}")
    backend = FakeBackend()
    before = metadata.read_bytes()
    with pytest.raises(ConfigurationError, match="paths alias"):
        annotate_metadata(
            metadata_csv=metadata,
            audio_dir=audio_dir,
            output_csv=alias,
            backend=backend,
            options=AnnotationOptions(overwrite=True),
        )
    assert backend.batches == []
    assert metadata.read_bytes() == before


def test_annotation_rejects_default_diagnostics_aliasing_source(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    aliased_metadata = tmp_path / "out.diagnostics.json"
    metadata.replace(aliased_metadata)
    backend = FakeBackend()
    before = aliased_metadata.read_bytes()
    with pytest.raises(ConfigurationError, match="paths alias"):
        annotate_metadata(
            metadata_csv=aliased_metadata,
            audio_dir=audio_dir,
            output_csv=tmp_path / "out.csv",
            backend=backend,
            options=AnnotationOptions(overwrite=True),
        )
    assert backend.batches == []
    assert aliased_metadata.read_bytes() == before


def test_failure_writer_overwrites_and_translates_oserror(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    failure = AlignmentFailure("r", "bad", "test", "failed")
    path = tmp_path / "failures.json"
    path.write_text("old", encoding="utf-8")
    _write_failures((failure,), path, overwrite=True)
    assert json.loads(path.read_text(encoding="utf-8"))[0]["code"] == "bad"

    def fail_replace(source: Path, destination: Path) -> None:
        raise OSError(source, destination)

    monkeypatch.setattr(os, "replace", fail_replace)
    with pytest.raises(BackendOutputError, match="publish"):
        _write_failures((failure,), tmp_path / "new.json", overwrite=False)


def test_duplicate_line_ids_remain_separate_ordered_rows(
    metadata_fixture: tuple[Path, Path], tmp_path: Path
) -> None:
    metadata, audio_dir = metadata_fixture
    run = annotate_metadata(
        metadata_csv=metadata,
        audio_dir=audio_dir,
        output_csv=tmp_path / "out.csv",
        backend=FakeBackend(),
    )
    assert [outcome.request_id for outcome in run.outcomes] == ["000000", "000001"]