simbarag/tests/unit/test_job_reliability.py

"""Real SQLite ORM delivery claims; providers/runtime are replaced locally."""

import asyncio
import os
import sys
from datetime import timedelta
from unittest.mock import AsyncMock

import pytest
from quart import Quart
from tortoise import Tortoise

from blueprints.cron.models import CronJob
from blueprints.cron.scheduler import execute_job
from blueprints.scheduled_messages import delivery
from blueprints.scheduled_messages.models import (
    JobAttempt,
    JobDelivery,
    MessageChannel,
    MessageStatus,
    Recurrence,
    ScheduledMessage,
)
from blueprints.scheduled_messages.scheduler import (
    _schedule_next_occurrence,
    process_due_messages,
    process_message,
    run_id,
)
from blueprints.users.models import User

MODULES = [
    "blueprints.users.models",
    "blueprints.scheduled_messages.models",
    "blueprints.cron.models",
]


@pytest.fixture
async def db(tmp_path):
    path = tmp_path / "jobs.sqlite3"
    await Tortoise.init(db_url=f"sqlite://{path}", modules={"models": MODULES})
    await Tortoise.generate_schemas()
    yield path
    await Tortoise.close_connections()


@pytest.fixture
async def owner(db):
    return await User.create(username="owner", email="owner@example.com")


@pytest.fixture
async def message(owner):
    return await ScheduledMessage.create(
        recipient="recipient",
        channel=MessageChannel.EMAIL,
        subject="subject",
        content="prompt",
        scheduled_at=delivery.now() - timedelta(minutes=1),
        created_by=owner,
    )


@pytest.fixture
async def job(owner):
    return await CronJob.create(
        name="cron",
        prompt="prompt",
        recipient="recipient",
        channel=MessageChannel.EMAIL,
        created_by=owner,
    )


@pytest.fixture
async def providers(monkeypatch):
    send = AsyncMock()
    agent = AsyncMock(return_value="response")
    monkeypatch.setattr(delivery, "send", send)
    monkeypatch.setattr(delivery, "agent_content", agent)
    return send, agent


async def test_atomic_concurrent_scheduled_claims(message, providers):
    results = await asyncio.gather(*[process_message(message) for _ in range(8)])
    assert results.count("sent") == 1
    assert providers[0].await_count == 1
    assert await JobAttempt.all().count() == 1
    await message.refresh_from_db()
    assert message.status == MessageStatus.SENT


async def test_claims_across_independent_processes(db):
    # Both processes use their own SQLite connection and ORM registry.
    code = """
import asyncio, sys
from tortoise import Tortoise
from blueprints.scheduled_messages.delivery import claim
async def main():
    await Tortoise.init(db_url="sqlite://" + sys.argv[1], modules={"models": ["blueprints.users.models", "blueprints.scheduled_messages.models", "blueprints.cron.models"]})
    result = await claim("cron", "same-job", "same-run")
    print("won" if result else "lost")
    await Tortoise.close_connections()
asyncio.run(main())
"""

    async def worker():
        p = await asyncio.create_subprocess_exec(
            sys.executable,
            "-c",
            code,
            str(db),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=os.environ.copy(),
        )
        out, err = await p.communicate()
        assert p.returncode == 0, err.decode()
        return out.decode().strip()

    assert sorted(await asyncio.gather(worker(), worker())) == ["lost", "won"]
    assert await JobAttempt.all().count() == 1


async def test_safe_send_failure_explicit_retry(message, providers):
    providers[0].side_effect = delivery.SafeDeliveryError("rejected before send")
    assert await process_message(message) == "failed"
    assert (await process_due_messages())["sent"] == 0
    providers[0].side_effect = None
    await message.refresh_from_db()
    assert await process_message(message, retry=True) == "sent"
    assert [a.status for a in await JobAttempt.all().order_by("started_at")] == [
        "failed",
        "sent",
    ]
    assert providers[0].await_count == 2


async def test_ambiguous_provider_failure_never_resends(message, providers):
    providers[0].side_effect = TimeoutError("may have delivered")
    assert await process_message(message) == "uncertain"
    await message.refresh_from_db()
    assert await process_message(message, retry=True) == "skipped"
    await process_due_messages()
    assert providers[0].await_count == 1
    assert (await JobAttempt.first()).status == "uncertain"


@pytest.mark.parametrize("phase", ["processing", "sending"])
async def test_abandoned_claim_is_uncertain(message, providers, phase):
    claimed = await delivery.claim("scheduled", message.id, run_id(message))
    if phase == "sending":
        await delivery.transition(claimed, "sending")
    await JobDelivery.filter(id=claimed.id).update(
        claimed_at=delivery.now() - timedelta(hours=2)
    )
    await ScheduledMessage.filter(id=message.id).update(status=MessageStatus.PROCESSING)
    await process_due_messages()
    await message.refresh_from_db()
    assert message.status == MessageStatus.UNCERTAIN
    assert (await JobAttempt.first()).status == "uncertain"
    providers[0].assert_not_awaited()


async def test_crash_after_provider_send_never_resends(message, providers):
    class WorkerCrash(BaseException):
        pass

    async def delivered_then_crash(*args):
        raise WorkerCrash()

    providers[0].side_effect = delivered_then_crash
    with pytest.raises(WorkerCrash):
        await process_message(message)
    assert (await JobDelivery.first()).status == "sending"
    await JobDelivery.all().update(claimed_at=delivery.now() - timedelta(hours=2))
    await process_due_messages()
    await message.refresh_from_db()
    assert message.status == MessageStatus.UNCERTAIN
    assert providers[0].await_count == 1


async def test_agent_failure_safe_retry_uses_owner(message, providers):
    message.use_agent = True
    await message.save()
    providers[1].side_effect = ValueError("no text")
    assert await process_message(message) == "failed"
    providers[0].assert_not_awaited()
    providers[1].side_effect = None
    await message.refresh_from_db()
    assert await process_message(message, retry=True) == "sent"
    args = providers[1].await_args.args
    assert args[0].id == message.created_by_id
    assert args[-1] == "scheduled"


async def test_recurrence_deduplicates_and_skips_missed_occurrences(message, providers):
    message.recurrence = Recurrence.DAILY
    message.scheduled_at = delivery.now() - timedelta(days=3)
    await message.save()
    assert await process_message(message) == "sent"
    await asyncio.gather(*[_schedule_next_occurrence(message) for _ in range(5)])
    assert await ScheduledMessage.all().count() == 2
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    assert successor.scheduled_at > delivery.now()
    assert successor.series_id == message.id
    await process_due_messages()
    assert await ScheduledMessage.all().count() == 2


async def test_cancelled_series_does_not_regenerate(message, providers):
    message.recurrence = Recurrence.DAILY
    await message.save()
    await process_message(message)
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    await ScheduledMessage.filter(id=successor.id).update(
        status=MessageStatus.CANCELLED
    )
    await _schedule_next_occurrence(message)
    await process_due_messages()
    assert await ScheduledMessage.all().count() == 2
    assert providers[0].await_count == 1


async def test_legacy_sent_rows_are_not_forked_into_new_series(owner, providers):
    # Rows sent before delivery claims existed have no successor links or
    # delivery records; the repair pass must leave them alone.
    common = dict(
        recipient="recipient",
        channel=MessageChannel.EMAIL,
        subject="subject",
        content="prompt",
        recurrence=Recurrence.DAILY,
        created_by=owner,
    )
    for days in (2, 1):
        await ScheduledMessage.create(
            **common,
            status=MessageStatus.SENT,
            scheduled_at=delivery.now() - timedelta(days=days),
        )
    await ScheduledMessage.create(
        **common, scheduled_at=delivery.now() + timedelta(hours=12)
    )
    await process_due_messages()
    assert await ScheduledMessage.filter(status=MessageStatus.PENDING).count() == 1
    assert await ScheduledMessage.all().count() == 3
    providers[0].assert_not_awaited()


async def test_claimed_send_without_successor_is_repaired(message, providers):
    message.recurrence = Recurrence.DAILY
    await message.save()
    assert await process_message(message) == "sent"
    # Simulate a crash after the send committed but before the successor existed.
    await ScheduledMessage.filter(previous_occurrence_id=message.id).delete()
    await process_due_messages()
    assert await ScheduledMessage.filter(previous_occurrence_id=message.id).exists()


async def test_cancellation_during_agent_prevents_delivery(message, providers):
    message.use_agent = True
    await message.save()

    async def cancel(*args):
        await ScheduledMessage.filter(id=message.id).update(
            status=MessageStatus.CANCELLED
        )
        return "response"

    providers[1].side_effect = cancel
    assert await process_message(message) == "cancelled"
    providers[0].assert_not_awaited()
    assert (await JobDelivery.first()).status == "cancelled"


async def test_duplicate_cron_triggers_scoped_to_run(job, providers):
    results = await asyncio.gather(
        *[execute_job(job, "occurrence-1") for _ in range(6)]
    )
    assert sum(bool(r.get("ok")) for r in results) == 1
    assert providers[0].await_count == 1
    assert await execute_job(job, "occurrence-2") == {
        "ok": True,
        "status": "sent",
        "error": None,
    }
    assert await JobAttempt.all().count() == 2
    assert providers[1].await_args.args[0].id == job.created_by_id
    assert providers[1].await_args.args[-1] == "cron"


async def test_legacy_cron_has_no_implicit_admin(job, providers):
    job.created_by_id = None
    await job.save()
    result = await execute_job(job, "legacy")
    assert result["status"] == "failed"
    assert "explicit" in result["error"]
    providers[1].assert_not_awaited()
    providers[0].assert_not_awaited()


async def test_cron_failure_retry_and_ambiguous_rejection(job, providers):
    providers[1].side_effect = ValueError("agent failed")
    assert (await execute_job(job, "run"))["status"] == "failed"
    assert (await execute_job(job, "run"))["skipped"]
    providers[1].side_effect = None
    providers[0].side_effect = TimeoutError("ambiguous")
    assert (await execute_job(job, "run", retry=True))["status"] == "uncertain"
    assert (await execute_job(job, "run", retry=True))["skipped"]
    assert await JobAttempt.all().count() == 2
    assert providers[0].await_count == 1


async def test_disabled_cron_prevents_send_after_agent(job, providers):
    async def disable(*args):
        await CronJob.filter(id=job.id).update(enabled=False)
        return "response"

    providers[1].side_effect = disable
    assert (await execute_job(job, "run"))["status"] == "cancelled"
    providers[0].assert_not_awaited()


async def test_history_records_attempts(message, providers):
    await process_message(message)
    result = await delivery.history("scheduled", message.id)
    assert result[0]["status"] == "sent"
    assert not result[0]["retryable"]
    assert result[0]["attempts"][0]["finished_at"]


@pytest.fixture
async def api(owner, monkeypatch):
    from quart_jwt_extended import JWTManager, create_access_token
    from blueprints.cron import cron_blueprint
    from blueprints.scheduled_messages import scheduled_messages_blueprint

    app = Quart(__name__)
    app.config["JWT_SECRET_KEY"] = "test-secret" * 8
    JWTManager(app)
    app.register_blueprint(cron_blueprint)
    app.register_blueprint(scheduled_messages_blueprint)
    monkeypatch.setenv("CRON_SECRET", "cron-secret")
    owner.ldap_groups = ["lldap_admin"]
    await owner.save()
    async with app.app_context():
        token = create_access_token(identity=str(owner.id))
    return app.test_client(), {"Authorization": f"Bearer {token}"}


async def test_trigger_auth_and_required_cron_identity(job, api, providers):
    client, _ = api
    url = f"/api/cron/{job.id}/trigger"
    assert (await client.post(url, json={"run_id": "run"})).status_code == 401
    assert (
        await client.post(url, headers={"X-Cron-Secret": "cron-secret"})
    ).status_code == 400
    headers = {"X-Cron-Secret": "cron-secret", "Idempotency-Key": "run"}
    assert (await client.post(url, headers=headers)).status_code == 200
    assert (await (await client.post(url, headers=headers)).get_json())["skipped"]
    assert providers[0].await_count == 1


async def test_admin_history_retry_and_failed_reschedule(message, providers, api):
    client, headers = api
    url = f"/api/scheduled-messages/{message.id}"
    providers[0].side_effect = delivery.SafeDeliveryError("not delivered")
    await process_message(message)
    response = await client.put(
        url,
        headers=headers,
        json={"scheduled_at": (delivery.now() + timedelta(hours=1)).isoformat()},
    )
    assert response.status_code == 200
    assert (await response.get_json())["status"] == "failed"
    assert (await client.post(url + "/retry")).status_code == 401
    providers[0].side_effect = None
    assert (await client.post(url + "/retry", headers=headers)).status_code == 200
    data = await (await client.get(url + "/deliveries", headers=headers)).get_json()
    assert len(data[0]["attempts"]) == 2


async def test_api_rejects_ambiguous_retry(message, providers, api):
    client, headers = api
    providers[0].side_effect = TimeoutError("unknown")
    await process_message(message)
    assert (
        await client.post(
            f"/api/scheduled-messages/{message.id}/retry", headers=headers
        )
    ).status_code == 409
    assert providers[0].await_count == 1


async def test_api_cancel_recurring_series(message, providers, api):
    client, headers = api
    message.recurrence = Recurrence.DAILY
    await message.save()
    await process_message(message)
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    response = await client.put(
        f"/api/scheduled-messages/{message.id}",
        headers=headers,
        json={"status": "cancelled"},
    )
    assert response.status_code == 200
    await successor.refresh_from_db()
    assert successor.status == MessageStatus.CANCELLED
    await process_due_messages()
    assert await ScheduledMessage.all().count() == 2


async def test_cancel_keeps_delivered_status_and_stops_series(message, providers, api):
    client, headers = api
    message.recurrence = Recurrence.DAILY
    await message.save()
    await process_message(message)
    response = await client.put(
        f"/api/scheduled-messages/{message.id}",
        headers=headers,
        json={"status": "cancelled"},
    )
    assert response.status_code == 200
    await message.refresh_from_db()
    assert message.status == MessageStatus.SENT  # delivered stays delivered
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    assert successor.status == MessageStatus.CANCELLED


async def test_cancel_during_send_does_not_claim_cancelled(message, providers, api):
    client, headers = api
    message.recurrence = Recurrence.DAILY
    await message.save()
    in_flight, release = asyncio.Event(), asyncio.Event()

    async def slow_send(*args):
        in_flight.set()
        await release.wait()

    providers[0].side_effect = slow_send
    task = asyncio.create_task(process_message(message))
    await in_flight.wait()
    response = await client.put(
        f"/api/scheduled-messages/{message.id}",
        headers=headers,
        json={"status": "cancelled"},
    )
    assert response.status_code == 200
    release.set()
    assert await task == "sent"
    await message.refresh_from_db()
    assert message.status == MessageStatus.SENT
    # The series stops even though nothing pending existed when cancelled.
    await process_due_messages()
    assert not await ScheduledMessage.filter(
        status=MessageStatus.PENDING, deleted_at__isnull=True
    ).exists()


async def test_deleting_edited_occurrence_does_not_revive_series(
    message, providers, api
):
    client, headers = api
    message.recurrence = Recurrence.DAILY
    await message.save()
    await process_message(message)
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    url = f"/api/scheduled-messages/{successor.id}"
    assert (
        await client.put(url, headers=headers, json={"recurrence": "none"})
    ).status_code == 200
    assert (await client.delete(url, headers=headers)).status_code == 200
    await process_due_messages()
    assert not await ScheduledMessage.filter(
        status=MessageStatus.PENDING, deleted_at__isnull=True
    ).exists()


async def test_api_assign_legacy_owner(job, api, providers, owner):
    job.created_by_id = None
    await job.save()
    client, headers = api
    response = await client.put(
        f"/api/cron/{job.id}", headers=headers, json={"created_by": str(owner.id)}
    )
    assert response.status_code == 200
    await job.refresh_from_db()
    assert (await execute_job(job, "run"))["ok"]


async def test_due_snapshot_does_not_send_rescheduled_message(message, providers):
    await ScheduledMessage.filter(id=message.id).update(
        scheduled_at=delivery.now() + timedelta(hours=1), content="edited"
    )
    assert await process_message(message) == "skipped"
    providers[0].assert_not_awaited()


async def test_claim_loads_current_content(message, providers):
    await ScheduledMessage.filter(id=message.id).update(content="edited")
    assert await process_message(message) == "sent"
    assert providers[0].await_args.args[2] == "edited"


async def test_delete_recurring_hides_tombstone_and_prevents_repair(
    message, providers, api
):
    message.recurrence = Recurrence.DAILY
    await message.save()
    await process_message(message)
    successor = await ScheduledMessage.get(previous_occurrence_id=message.id)
    client, headers = api
    response = await client.delete(
        f"/api/scheduled-messages/{successor.id}", headers=headers
    )
    assert response.status_code == 200
    rows = await (
        await client.get("/api/scheduled-messages/", headers=headers)
    ).get_json()
    assert [r["id"] for r in rows] == [str(message.id)]
    await process_due_messages()
    assert await ScheduledMessage.all().count() == 2
    assert (await ScheduledMessage.get(id=successor.id)).deleted_at is not None


@pytest.mark.parametrize(
    "status, safe", [(200, None), (401, True), (429, True), (500, False)]
)
async def test_email_provider_acknowledgement(monkeypatch, status, safe):
    import httpx

    monkeypatch.setenv("MAILGUN_API_KEY", "test")
    monkeypatch.setenv("MAILGUN_DOMAIN", "test.invalid")
    transport = httpx.MockTransport(
        lambda request: httpx.Response(status, json={"id": "test"})
    )
    original = httpx.AsyncClient
    monkeypatch.setattr(httpx, "AsyncClient", lambda: original(transport=transport))
    if safe is None:
        await delivery.send(MessageChannel.EMAIL, "recipient", "body", "subject")
    else:
        with pytest.raises(
            delivery.SafeDeliveryError if safe else httpx.HTTPStatusError
        ):
            await delivery.send(MessageChannel.EMAIL, "recipient", "body", "subject")


async def test_missing_email_configuration_is_safe_failure(monkeypatch):
    monkeypatch.delenv("MAILGUN_API_KEY", raising=False)
    monkeypatch.delenv("MAILGUN_DOMAIN", raising=False)
    with pytest.raises(delivery.SafeDeliveryError):
        await delivery.send(MessageChannel.EMAIL, "recipient", "body")


async def test_agent_adapter_uses_shared_runtime(owner, monkeypatch):
    import types

    runtime = types.ModuleType("blueprints.conversation.runtime")
    invoke = AsyncMock(return_value={"messages": [{"content": "answer"}]})
    runtime.invoke_agent_for_user = invoke
    monkeypatch.setitem(sys.modules, "blueprints.conversation.runtime", runtime)
    assert (
        await delivery.agent_content(owner, "query", "thread", "scheduled") == "answer"
    )
    invoke.assert_awaited_once_with(owner, "query", "thread", channel="scheduled")


async def test_migration_snapshot_preserves_prior_models_and_describes_new_models(db):
    import importlib.util
    from pathlib import Path
    from aerich.utils import compress_dict, decompress_dict, get_models_describe

    def load(path):
        spec = importlib.util.spec_from_file_location("test_migration", path)
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module

    migration_dir = Path(__file__).resolve().parents[2] / "migrations" / "models"
    previous = load(migration_dir / "19_20261012225640_add_fk_indexes.py")
    migration = load(migration_dir / "20_20261012231017_application_foundation.py")
    state = decompress_dict(migration.MODELS_STATE)
    owned = {
        "models." + name
        for name in ("ScheduledMessage", "CronJob", "JobDelivery", "JobAttempt")
    }
    current = decompress_dict(compress_dict(get_models_describe("models")))
    for name in owned:
        assert state[name] == current[name]
    assert "models.User" in decompress_dict(previous.MODELS_STATE)
    sql = await migration.upgrade(None)
    assert 'UNIQUE ("kind", "job_id", "run_id")' in sql
    assert '"previous_occurrence_id" UUID UNIQUE' in sql
    assert '"created_by_id" UUID' in sql


@pytest.mark.parametrize(
    "content, expected",
    [
        ([{"type": "text", "text": "answer"}], "answer"),
        ("<think>private reasoning</think>answer", "answer"),
    ],
)
async def test_agent_adapter_extracts_supported_content(
    owner, monkeypatch, content, expected
):
    import types

    runtime = types.ModuleType("blueprints.conversation.runtime")
    runtime.invoke_agent_for_user = AsyncMock(
        return_value={"messages": [{"content": content}]}
    )
    monkeypatch.setitem(sys.modules, "blueprints.conversation.runtime", runtime)
    assert await delivery.agent_content(owner, "query", "thread", "cron") == expected


@pytest.mark.parametrize(
    "path, body",
    [
        ("/api/scheduled-messages/", ["array"]),
        ("/api/scheduled-messages/", {"recipient": 5, "channel": "email"}),
        ("/api/cron/", ["array"]),
        ("/api/cron/", {"name": 7, "prompt": "p", "recipient": "r"}),
    ],
)
async def test_malformed_create_bodies_are_rejected(api, path, body):
    client, headers = api
    resp = await client.post(path, headers=headers, json=body)
    assert resp.status_code == 400


async def test_malformed_update_and_trigger_bodies(job, message, api, providers):
    client, headers = api
    for path in (f"/api/scheduled-messages/{message.id}", f"/api/cron/{job.id}"):
        resp = await client.put(path, headers=headers, json={"recipient": ["x"]})
        assert resp.status_code == 400
    resp = await client.post(
        f"/api/cron/{job.id}/trigger",
        headers={"X-Cron-Secret": "cron-secret"},
        json=["run"],
    )
    assert resp.status_code == 400


async def test_update_rejects_null_and_blank_required_fields(job, message, api):
    client, headers = api
    for path, field in (
        (f"/api/scheduled-messages/{message.id}", "recipient"),
        (f"/api/scheduled-messages/{message.id}", "content"),
        (f"/api/cron/{job.id}", "name"),
        (f"/api/cron/{job.id}", "prompt"),
    ):
        for value in (None, "   "):
            resp = await client.put(path, headers=headers, json={field: value})
            assert resp.status_code == 400, (path, field, value)
    # subject may be cleared, except where the channel requires it (email).
    resp = await client.put(
        f"/api/scheduled-messages/{message.id}", headers=headers, json={"subject": None}
    )
    assert resp.status_code == 400
    resp = await client.put(
        f"/api/scheduled-messages/{message.id}",
        headers=headers,
        json={"subject": None, "channel": "imessage"},
    )
    assert resp.status_code == 200


async def test_cancel_started_one_off_message_is_rejected(message, providers, api):
    client, headers = api
    in_flight, release = asyncio.Event(), asyncio.Event()

    async def slow_send(*args):
        in_flight.set()
        await release.wait()

    providers[0].side_effect = slow_send
    task = asyncio.create_task(process_message(message))
    await in_flight.wait()
    response = await client.put(
        f"/api/scheduled-messages/{message.id}",
        headers=headers,
        json={"status": "cancelled"},
    )
    assert response.status_code == 409
    release.set()
    assert await task == "sent"
    assert await ScheduledMessage.all().count() == 1  # no tombstone for one-offs