clickstream-ch-kafka-superset-demo/generator/tests/test_startup_history_artifact.py

"""
Тесты портативного артефакта стартовой истории.
"""

from dataclasses import replace
from datetime import datetime, timezone

import pytest

from generator import GeneratorState


def _state() -> GeneratorState:
    import random

    now = datetime(2026, 1, 1, 1, 0, tzinfo=timezone.utc)
    return GeneratorState(
        tick=1,
        rng_state=random.Random(42).getstate(),
        last_batch_id="startup-history-abc",
        last_timestamp=now,
        model_timestamp=now,
        wall_timestamp=now,
        model_time_speed=1,
        model_timezone="UTC",
        model_t0=datetime(2026, 1, 1, 0, 0, tzinfo=timezone.utc),
        gen_seed=42,
        population=[
            {
                "user_domain_id": "user-1",
                "seed_click_id": "seed-1",
                "active_click_id": None,
                "last_finished_at": None,
            }
        ],
    )


def _batch() -> dict[str, list[dict]]:
    return {
        "browser_events": [
            {
                "event_id": "event-1",
                "click_id": "click-1",
                "user_domain_id": "user-1",
                "event_timestamp": "2026-01-01 00:00:00.000000",
            }
        ],
        "location_events": [
            {
                "event_id": "event-1",
                "page_url_path": "/home",
            }
        ],
        "device_events": [
            {
                "click_id": "click-1",
                "user_domain_id": "user-1",
            }
        ],
        "geo_events": [
            {
                "click_id": "click-1",
                "country": "RU",
            }
        ],
    }


def test_artifact_roundtrip_keeps_events_state_and_manifest(base_config):
    """Артефакт хранит события, state и manifest одним проверяемым набором."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        load_startup_history_artifact,
        validate_startup_history_artifact,
        write_startup_history_artifact,
    )

    state = _state()
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    manifest = build_manifest(
        config=replace(
            base_config,
            model_t0=state.model_t0,
            model_t_end=state.model_timestamp,
            model_time_speed=1,
            model_timezone="UTC",
            seed=42,
        ),
        counters=builder.counters,
        state=state,
    )

    artifact = builder.to_artifact(manifest=manifest, state=state)
    (
        loaded_state,
        loaded_manifest,
        loaded_topics,
        loaded_raw_topics,
    ) = validate_startup_history_artifact(
        artifact,
        expected_config=replace(
            base_config,
            model_t0=state.model_t0,
            model_t_end=state.model_timestamp,
            model_time_speed=1,
            model_timezone="UTC",
            seed=42,
        ),
    )

    assert artifact["artifact_version"] == "1.0"
    assert loaded_state.last_batch_id == state.last_batch_id
    assert loaded_manifest["state"]["last_batch_id"] == state.last_batch_id
    assert loaded_topics["browser_events"][0]["event_id"] == "event-1"
    assert loaded_raw_topics["browser_events"][0]["value_json"].encode("utf-8") == (
        b'{"event_id": "event-1", "click_id": "click-1", '
        b'"user_domain_id": "user-1", '
        b'"event_timestamp": "2026-01-01 00:00:00.000000"}'
    )

    path = base_config.data_dir.parent / "tmp-startup-history-artifact.json"
    try:
        write_startup_history_artifact(path, artifact)
        loaded = load_startup_history_artifact(path)
        assert loaded["manifest"]["totals"]["events"] == 1
    finally:
        path.unlink(missing_ok=True)


def test_manifest_and_artifact_show_launch_profile(base_config):
    """Manifest и файл артефакта показывают выбранный профиль запуска."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        load_startup_history_artifact,
        write_startup_history_artifact,
    )

    state = _state()
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    manifest = build_manifest(
        config=replace(
            base_config,
            model_t0=state.model_t0,
            model_t_end=state.model_timestamp,
            model_time_speed=1,
            model_timezone="UTC",
            seed=42,
            launch_profile="daily-wave",
        ),
        counters=builder.counters,
        state=state,
    )
    artifact = builder.to_artifact(manifest=manifest, state=state)

    path = base_config.data_dir.parent / "tmp-startup-history-profile.json"
    try:
        write_startup_history_artifact(path, artifact)
        loaded = load_startup_history_artifact(path)
        assert manifest["launch_profile"] == "daily-wave"
        assert loaded["manifest"]["launch_profile"] == "daily-wave"
    finally:
        path.unlink(missing_ok=True)


def test_artifact_validation_rejects_mismatched_manifest(base_config):
    """Валидация отвергает артефакт, где manifest не совпадает с событиями."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        validate_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    manifest = build_manifest(config=config, counters=builder.counters, state=state)
    artifact = builder.to_artifact(manifest=manifest, state=state)
    artifact["manifest"]["gen_seed"] = 43

    with pytest.raises(ValueError, match="GEN_SEED"):
        validate_startup_history_artifact(artifact, expected_config=config)


def test_manifest_clickhouse_stats_detect_mismatch():
    """Сверка ClickHouse с manifest ловит расхождение контрольных чисел."""
    from clickstream_generator.startup_history_artifact import (
        compare_clickhouse_stats_to_manifest,
    )

    manifest = {
        "totals": {
            "events": 10,
            "visits": 4,
            "users": 3,
            "min_event_timestamp": "2026-01-01 00:00:00.000000",
            "max_event_timestamp": "2026-01-01 00:10:00.000000",
        }
    }
    stats = {
        "events": "9",
        "visits": "4",
        "users": "3",
        "min_event_timestamp": "2026-01-01 00:00:00.000000",
        "max_event_timestamp": "2026-01-01 00:10:00.000000",
    }

    assert compare_clickhouse_stats_to_manifest(manifest, stats) == ["events"]


def test_raw_topic_value_can_keep_original_json_bytes(base_config):
    """Raw value проверяется как JSON, но импортируется без пересборки строки."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )
    raw_value = (
        '{"event_timestamp":"2026-01-01 00:00:00.000000",'
        '"user_domain_id":"user-1","click_id":"click-1","event_id":"event-1"}'
    )
    artifact["raw_topics"]["browser_events"][0]["value_json"] = raw_value

    class Publisher:
        def __init__(self):
            self.records = None

        def publish_records(self, topic, records):
            if topic == "browser_events":
                self.records = records
            return len(records), 0

        def flush(self):
            return None

    class CompactWriter:
        def save(self, value):
            self.saved = value

        def flush(self):
            return None

    publisher = Publisher()
    import_startup_history_artifact(
        artifact,
        publisher=publisher,
        state_manager=CompactWriter(),
        manifest_manager=CompactWriter(),
        expected_config=config,
    )

    assert publisher.records[0]["value_json"] == raw_value


def test_artifact_v1_requires_raw_topics_for_import(base_config):
    """Artifact v1.0 без raw value не может обещать byte-for-byte импорт."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )
    del artifact["raw_topics"]

    class Publisher:
        def publish_records(self, topic, records):
            raise AssertionError("publish must not be called")

    with pytest.raises(ValueError, match="raw_topics"):
        import_startup_history_artifact(
            artifact,
            publisher=Publisher(),
            state_manager=None,
            manifest_manager=None,
            expected_config=config,
        )


def test_import_rejects_publisher_without_raw_publish_records(base_config):
    """Импорт не должен пересобирать dict, если publisher не умеет raw records."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )
    artifact["raw_topics"]["browser_events"][0]["value_json"] = (
        '{"event_timestamp":"2026-01-01 00:00:00.000000",'
        '"user_domain_id":"user-1","click_id":"click-1","event_id":"event-1"}'
    )

    class DictPublisher:
        def publish(self, topic, events):
            raise AssertionError("dict publish must not be called")

        def flush(self):
            return None

    with pytest.raises(TypeError, match="publish_records"):
        import_startup_history_artifact(
            artifact,
            publisher=DictPublisher(),
            state_manager=None,
            manifest_manager=None,
            expected_config=config,
        )


def test_import_rejects_dirty_kafka_topics_before_publish(base_config):
    """Повторный импорт не должен дописывать дубли в непустые data-топики."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )

    class Publisher:
        def publish(self, topic, events):
            raise AssertionError("publish must not be called")

    class DirtyTopics:
        def assert_data_topics_empty(self):
            raise RuntimeError("Kafka data topics are not empty")

    with pytest.raises(RuntimeError, match="not empty"):
        import_startup_history_artifact(
            artifact,
            publisher=Publisher(),
            state_manager=None,
            manifest_manager=None,
            expected_config=config,
            topic_inspector=DirtyTopics(),
        )


def test_import_rolls_back_kafka_topics_after_partial_publish(base_config):
    """При частичной публикации импорт откатывает Kafka-топики до повтора."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )

    class PartialPublisher:
        def __init__(self):
            self.calls = []

        def publish_records(self, topic, records):
            self.calls.append(topic)
            if topic == "location_events":
                raise RuntimeError("location down")
            return len(records), 0

    class CompactWriter:
        def save(self, value):
            raise AssertionError("compact topics must not be written")

    class Inspector:
        def __init__(self):
            self.snapshot = None
            self.rollback = None

        def assert_data_topics_empty(self):
            return None

        def snapshot_import_topics(self):
            self.snapshot = {"topics": ["browser_events", "location_events"]}
            return self.snapshot

        def rollback_import_topics(self, snapshot):
            self.rollback = snapshot

    publisher = PartialPublisher()
    inspector = Inspector()

    with pytest.raises(RuntimeError, match="location down"):
        import_startup_history_artifact(
            artifact,
            publisher=publisher,
            state_manager=CompactWriter(),
            manifest_manager=CompactWriter(),
            expected_config=config,
            topic_inspector=inspector,
        )

    assert publisher.calls == ["browser_events", "location_events"]
    assert inspector.rollback == inspector.snapshot


def test_import_replays_events_and_compact_topics(base_config):
    """Импорт пишет события в Kafka и служебные compact-топики, не трогая ClickHouse."""
    from clickstream_generator.startup_history_artifact import (
        StartupHistoryArtifactBuilder,
        build_manifest,
        import_startup_history_artifact,
    )

    state = _state()
    config = replace(
        base_config,
        model_t0=state.model_t0,
        model_t_end=state.model_timestamp,
        model_time_speed=1,
        model_timezone="UTC",
        seed=42,
    )
    builder = StartupHistoryArtifactBuilder()
    builder.add_batch(_batch())
    artifact = builder.to_artifact(
        manifest=build_manifest(config=config, counters=builder.counters, state=state),
        state=state,
    )

    class Publisher:
        def __init__(self):
            self.calls = []
            self.flushed = False

        def publish_records(self, topic, records):
            self.calls.append((topic, records))
            return len(records), 0

        def flush(self):
            self.flushed = True

    class CompactWriter:
        def __init__(self):
            self.saved = None
            self.flushed = False

        def save(self, value):
            self.saved = value

        def flush(self):
            self.flushed = True

    publisher = Publisher()
    state_manager = CompactWriter()
    manifest_manager = CompactWriter()

    result = import_startup_history_artifact(
        artifact,
        publisher=publisher,
        state_manager=state_manager,
        manifest_manager=manifest_manager,
        expected_config=config,
    )

    assert [topic for topic, _ in publisher.calls] == [
        "browser_events",
        "location_events",
        "device_events",
        "geo_events",
    ]
    assert publisher.calls[0][1][0]["value_json"].encode("utf-8") == (
        b'{"event_id": "event-1", "click_id": "click-1", '
        b'"user_domain_id": "user-1", '
        b'"event_timestamp": "2026-01-01 00:00:00.000000"}'
    )
    assert result["events"] == 4
    assert state_manager.saved.last_batch_id == state.last_batch_id
    assert manifest_manager.saved["state"]["last_batch_id"] == state.last_batch_id
    assert publisher.flushed
    assert state_manager.flushed
    assert manifest_manager.flushed