[Bug] A recovered Batch job past its window with no execution metadata never finishes expiring #3074

import os
from datetime import datetime, timedelta, timezone

import pytest

os.environ.setdefault("SECRET_KEY", "test-secret-key-for-testing")

from aibrix.batch.batch_manager import BatchManager
from aibrix.batch.job_entity import (
    BatchJob,
    BatchJobEndpoint,
    BatchJobSpec,
    BatchJobState,
    BatchJobStatus,
    CompletionWindow,
    ConditionType,
    ObjectMeta,
    TypeMeta,
)
from aibrix.context import InfrastructureContext
from aibrix.metadata.api.v1.batch import _batch_job_to_openai_response


class _Scheduler:
    def __init__(self):
        self.appended = []

    def append_job(self, job_id):
        self.appended.append(job_id)


def _overdue_in_progress_job(job_id, execution):
    created = datetime.now(timezone.utc) - timedelta(hours=25)
    job = BatchJob(
        typeMeta=TypeMeta(apiVersion="batch/v1", kind="Job"),
        metadata=ObjectMeta(
            name=f"job-{job_id}", namespace="default", uid=f"uid-{job_id}",
            creationTimestamp=created, resourceVersion=None, deletionTimestamp=None,
        ),
        spec=BatchJobSpec(
            input_file_id=f"input-{job_id}",
            endpoint=BatchJobEndpoint.CHAT_COMPLETIONS.value,
            completion_window=CompletionWindow.TWENTY_FOUR_HOURS.expires_at(),
        ),
        status=BatchJobStatus(
            jobID=job_id, state=BatchJobState.IN_PROGRESS,
            createdAt=created, inProgressAt=created,
        ),
    )
    job.status.request_counts.total = 1
    job.status.execution = execution
    return job


@pytest.mark.asyncio
@pytest.mark.parametrize("execution", [None, {}], ids=["execution-none", "execution-empty"])
async def test_recovered_overdue_in_progress_job_expires(execution):
    manager = BatchManager(InfrastructureContext())
    scheduler = _Scheduler()
    manager.set_scheduler(scheduler)
    job = _overdue_in_progress_job(f"overdue-{execution is None}", execution)

    await manager.job_committed_handler(job)  # recovery after a restart
    await manager.expire_job(job.job_id)  # the scheduler's next expiry pass

    assert scheduler.appended == []  # overdue work is not scheduled
    current = manager._done_jobs.get(job.job_id) or await manager.get_job(job.job_id)
    assert current.status.state == BatchJobState.FINALIZED, (
        f"state={current.status.state.value} condition={current.status.condition} "
        f"client={_batch_job_to_openai_response(current).status}"
    )
    assert current.status.condition == ConditionType.EXPIRED
    assert _batch_job_to_openai_response(current).status == "expired"