IPSExtractor/services/bot_descarga/tests/test_store.py

from __future__ import annotations

import importlib
import sys
from pathlib import Path

import pytest

from infrastructure.storage_minio import downloads_prefix, excel_key, pdf_key


def test_minio_key_helpers():
    assert excel_key("abc", "input.xlsx") == "uploads/abc/input.xlsx"
    assert downloads_prefix("abc") == "downloads/abc/"
    assert pdf_key("abc", "900", "123") == "downloads/abc/900/123.pdf"


@pytest.fixture()
def store(tmp_path, monkeypatch):
    db_url = f"sqlite:///{(tmp_path / 'bot.db').as_posix()}"
    monkeypatch.setenv("DATABASE_URL", db_url)

    # Recargar módulos para que el engine use SQLITE
    for name in [
        "infrastructure.database",
        "infrastructure.store",
    ]:
        if name in sys.modules:
            importlib.reload(sys.modules[name])
        else:
            importlib.import_module(name)

    from domain.models import JobStatus, WorkerStatus
    from infrastructure import database as db
    from infrastructure.store import JobStore

    importlib.reload(db)
    import infrastructure.store as store_mod

    importlib.reload(store_mod)
    return store_mod.JobStore(), JobStatus, WorkerStatus


def test_register_heartbeat_and_offline(store):
    job_store, _JobStatus, WorkerStatus = store
    w = job_store.register_worker("PC-01", hostname="host1")
    assert w.status == WorkerStatus.ONLINE.value
    assert job_store.is_worker_online("PC-01")

    job_store.heartbeat("PC-01")
    workers = job_store.list_workers()
    assert len(workers) == 1
    assert workers[0].id == "PC-01"


def test_create_job_requires_worker_and_claim_is_atomic(store):
    job_store, JobStatus, _WorkerStatus = store
    job_store.register_worker("PC-01")

    job = job_store.create_job(
        worker_id="PC-01",
        excel_key="uploads/x/input.xlsx",
        original_filename="datos.xlsx",
        title="Lote 1",
    )
    assert job.status == JobStatus.PENDING.value
    assert job.worker_id == "PC-01"
    assert job.downloads_prefix == downloads_prefix(job.id)

    claimed = job_store.claim_pending("PC-01")
    assert claimed is not None
    assert claimed.id == job.id
    assert claimed.status == JobStatus.RUNNING.value

    again = job_store.claim_pending("PC-01")
    assert again is None

    job_store.register_worker("PC-02")
    other = job_store.claim_pending("PC-02")
    assert other is None


def test_reregister_abandons_leftover_job(store):
    job_store, JobStatus, _WorkerStatus = store
    job_store.register_worker("PC-01")
    job = job_store.create_job(
        worker_id="PC-01",
        excel_key="uploads/x/input.xlsx",
        original_filename="datos.xlsx",
    )
    job_store.claim_pending("PC-01")

    job_store.register_worker("PC-01")
    leftover = job_store.get_job(job.id)
    assert leftover is not None
    assert leftover.status == JobStatus.CANCELLED.value
    assert job_store.claim_pending("PC-01") is None

    fresh = job_store.create_job(
        worker_id="PC-01",
        excel_key="uploads/y/input.xlsx",
        original_filename="nuevo.xlsx",
    )
    claimed = job_store.claim_pending("PC-01")
    assert claimed is not None
    assert claimed.id == fresh.id


def test_cancel_pending_marks_cancelled(store):
    job_store, JobStatus, _WorkerStatus = store
    job_store.register_worker("PC-01")
    job = job_store.create_job("PC-01", "uploads/j/input.xlsx", "a.xlsx")
    cancelled = job_store.request_cancel(job.id)
    assert cancelled is not None
    assert cancelled.status == JobStatus.CANCELLED.value
    assert cancelled.cancel_requested is True


def test_cancel_running_sets_flag(store):
    job_store, JobStatus, _WorkerStatus = store
    job_store.register_worker("PC-01")
    job = job_store.create_job("PC-01", "uploads/j/input.xlsx", "a.xlsx")
    job_store.claim_pending("PC-01")
    flagged = job_store.request_cancel(job.id)
    assert flagged is not None
    assert flagged.cancel_requested is True
    assert flagged.status == JobStatus.RUNNING.value
    assert job_store.is_cancel_requested(job.id) is True


def test_events_and_patch_complete(store):
    job_store, JobStatus, _WorkerStatus = store
    job_store.register_worker("PC-01")
    job = job_store.create_job("PC-01", "uploads/j/input.xlsx", "a.xlsx")
    job_store.claim_pending("PC-01")
    e1 = job_store.add_event(
        job.id, "radicado_saved", "ok", radicado="1", nit="9", index=1, total=2
    )
    assert e1.id >= 1
    events = job_store.list_events(job.id, since_id=0)
    assert len(events) == 1

    done = job_store.patch_job(
        job.id,
        status=JobStatus.COMPLETED.value,
        result={"exitosos": 1},
    )
    assert done is not None
    assert done.status == JobStatus.COMPLETED.value
    worker = job_store.get_worker("PC-01")
    assert worker is not None
    assert worker.current_job_id is None


def test_list_jobs_by_parent(store):
    job_store, _JobStatus, _WorkerStatus = store
    from infrastructure.database import Job as IpsJob
    from infrastructure.database import session_scope, utcnow

    job_store.register_worker("PC-01")
    with session_scope() as session:
        session.add(
            IpsJob(
                id="parent-1",
                status="pending",
                title="Trabajo",
                job_metadata={},
                created_at=utcnow(),
                updated_at=utcnow(),
            )
        )

    a = job_store.create_job("PC-01", "uploads/a/input.xlsx", "a.xlsx", job_id="parent-1")
    job_store.create_job("PC-01", "uploads/b/input.xlsx", "b.xlsx", job_id="other")
    listed = job_store.list_jobs(limit=50, job_id="parent-1")
    assert len(listed) == 1
    assert listed[0].id == a.id