testsuite/python/tests/test_172_5.py

############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Test jobcomp/kafka per-event default topics and enable flags."""

import json
import logging
import time

import pytest

import atf

kafka = pytest.importorskip("kafka")

logging.getLogger("kafka").setLevel(logging.WARNING)

# Only these events are enabled. The finish event is always on, and the
# others must stay silent whatever the job goes through.
enabled_events = ("hold", "requeue")
disabled_events = ("queue", "start", "release")
cluster = None


@pytest.fixture(scope="module", autouse=True)
def setup():
    global cluster
    atf.require_version(
        (26, 11),
        component="sbin/slurmctld",
        reason="Issue 51114: jobcomp job state transition events require 26.11+",
    )
    atf.require_build_config("HAVE_RDKAFKA")
    try:
        kafka.KafkaConsumer(bootstrap_servers=["localhost:9092"]).close()
    except kafka.errors.NoBrokersAvailable as e:
        pytest.skip(f"Kafka server not available at localhost:9092: {e}")
    atf.require_config_parameter("JobCompType", "jobcomp/kafka")
    # No topic names, so every event has to land on its default topic derived
    # from ClusterName, and only the enabled events may produce a record.
    atf.require_config_parameter(
        "JobCompParams", ",".join(f"enable_job_{event}" for event in enabled_events)
    )
    k_conf_path = atf.properties["slurm-config-dir"] + "/jobcomp_kafka_172_5.conf"
    atf.require_config_parameter("JobCompLoc", k_conf_path)
    atf.require_config_parameter_includes(
        "bootstrap.servers", "localhost:9092", source="jobcomp_kafka_172_5"
    )
    atf.require_config_parameter_includes(
        "enable.idempotence", "true", source="jobcomp_kafka_172_5"
    )
    # A requeued job is only eligible again after requeue_delay, which defaults
    # to the AuthInfo cred_expire (120 seconds), and is only started by the
    # next scheduling cycle. The (key, value) form replaces a value the site
    # config already sets.
    atf.require_config_parameter_includes("SchedulerParameters", ("requeue_delay", 0))
    atf.require_config_parameter_includes("SchedulerParameters", ("sched_interval", 1))
    atf.require_slurm_running()
    cluster = atf.get_config_parameter("ClusterName")


def topic_for(event):
    """Return the default topic of an event: ClusterName for finish."""
    if event == "finish":
        return cluster
    return f"{cluster}-job-{event}"


def utc_now():
    """Return the current UTC time in the format of the @event_time field."""
    return time.strftime("%FT%T", time.gmtime())


def read_records(topic_name, job_id, since):
    """Return the job's records on a topic produced at or after since.

    The default topics are shared with earlier runs, in which job ids are
    reused, so records older than the test are ignored. Every partition is
    read from its beginning up to the end offset taken when the read
    starts, so an empty poll on a slow broker cannot cut the read short.
    A topic nothing was produced to does not exist yet and has no records.
    """
    consumer = kafka.KafkaConsumer(bootstrap_servers=["localhost:9092"])
    records = []
    try:
        partitions = consumer.partitions_for_topic(topic_name)
        if not partitions:
            return records
        assignment = [kafka.TopicPartition(topic_name, p) for p in partitions]
        consumer.assign(assignment)
        consumer.seek_to_beginning()
        ends = consumer.end_offsets(assignment)
        for t in atf.timer(quiet=True, fatal=True):
            for batch in consumer.poll(timeout_ms=1000).values():
                for message in batch:
                    data = json.loads(message.value.decode("utf-8"))
                    if (
                        data.get("jobid") == job_id
                        and data.get("@event_time", "") >= since
                    ):
                        records.append(data)
            if all(consumer.position(tp) >= ends[tp] for tp in assignment):
                break
    finally:
        consumer.close()
    return records


def wait_for_records(event, job_id, since, count):
    """Return the job's records on the event's topic once there are count.

    Also checks that every record on that topic is of that event, so a
    record routed to the wrong topic fails here.
    """
    topic_name = topic_for(event)
    records = []
    for t in atf.timer():
        records = read_records(topic_name, job_id, since)
        if len(records) >= count:
            break
    assert (
        len(records) == count
    ), f"Expected {count} records of job {job_id} on {topic_name}, got {records}"
    for record in records:
        assert (
            record.get("event") == f"job {event}"
        ), f"Record of a different event on {topic_name}: {record}"
    return records


def assert_no_records(job_id, since, events):
    """Check that none of the events' topics has a record of the job."""
    for event in events:
        records = read_records(topic_for(event), job_id, since)
        assert (
            records == []
        ), f"Job {job_id} should have no record on {topic_for(event)}: {records}"


def test_hold_routed_to_its_topic():
    """A hold lands on the hold topic; the disabled release leaves no record."""
    since = utc_now()
    job_id = atf.submit_job_sbatch(
        "--begin=now+1hour -t1 -o /dev/null --wrap true", fatal=True
    )
    atf.wait_for_job_state(job_id, "PENDING", fatal=True)
    atf.run_command(f"scontrol hold {job_id}", fatal=True)
    atf.run_command(f"scontrol release {job_id}", fatal=True)
    atf.cancel_jobs([job_id], fatal=True)

    (finish,) = wait_for_records("finish", job_id, since, 1)
    assert finish["state"] == "CANCELLED", f"Finish record state: {finish['state']}"
    (hold,) = wait_for_records("hold", job_id, since, 1)
    assert hold["priority"] == 0, f"Hold record: {hold}"
    assert_no_records(job_id, since, disabled_events + ("requeue",))


def test_requeue_routed_to_its_topic():
    """A requeue lands on the requeue topic; start and queue leave none."""
    since = utc_now()
    job_id = atf.submit_job_sbatch(
        "--requeue -t5 -o /dev/null --wrap 'sleep infinity'", fatal=True
    )
    atf.wait_for_job_state(job_id, "RUNNING", fatal=True)
    atf.run_command(f"scontrol requeue {job_id}", fatal=True)
    atf.wait_for_job_state(job_id, "RUNNING", fatal=True)
    atf.cancel_jobs([job_id], fatal=True)

    finish1, finish2 = wait_for_records("finish", job_id, since, 2)
    assert (
        finish1["state"] == "REQUEUED"
    ), f"Finish record of the requeued run: {finish1['state']}"
    assert (
        finish2["state"] == "CANCELLED"
    ), f"Finish record of the second run: {finish2['state']}"
    (requeue,) = wait_for_records("requeue", job_id, since, 1)
    assert requeue["sluid"] == finish2["sluid"], f"Requeue record SLUID: {requeue}"
    assert_no_records(job_id, since, disabled_events + ("hold",))