testsuite/python/tests/test_157_4.py

############################################################################
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
############################################################################
"""Regression guardrails for bug 25211 job_submit/lua field clearing.

These exercise behavior that the fix must leave unchanged, so they are not
version gated and run on every supported version:

Group C - submit-path behavior is unchanged by the patch: assigning "" on
    submit leaves the field NULL, discarding the user-supplied value.

Group D - non-whitelist guardrails (must NOT regress).  Assigning "" to a
    field outside the whitelist must leave that field alone.  Every field
    here reaches a controller branch that would act on "": partition and
    qos revert to a default, mail_user reverts to the job owner, and
    std_out, std_err and std_in are cleared outright.  The plugin's
    whitelist is the only thing keeping "" away from those branches, and
    job_submit_plugins.shtml documents the fields as deliberately
    excluded.  Each case is seeded away from the value "" would produce,
    so a regression is observable rather than silently identical.  The
    Lua hook and the test are both generated from _GROUP_D_CASES.

Group E - assigning a non-empty value must still take effect.  job_desc is
    documented as input/output and the hook as able to modify the job
    parameters supplied by the user, so narrowing which fields "" clears
    must not disturb ordinary assignment:
    E1 features      set on modify
    E2 tres_per_node set on modify
    E3 comment       overwritten on modify
    E4 features      set on submit
    E5 layout        set on submit, read on submit, refused when malformed
                     (26.11 and later only)
"""

import os

import pytest

import atf

QOS_NAME = "test157_qos"
ACCOUNT_NAME = "test157_acct"
PARTITION_NAME = "test157_part"
LUA_LAYOUT = "lua157=set"
USER_LAYOUT = "user157=given"

requires_layout = pytest.mark.skipif(
    atf.get_version("bin/sbatch") < (26, 11)
    or atf.get_version("bin/scontrol") < (26, 11),
    reason="Issue 51191: the job layout was added in 26.11",
)

# ---------------------------------------------------------------------------
# Group D - fields outside the whitelist.  Both the Lua hook below and the
# Group D test are generated from this table, so a field cannot gain a hook
# branch without also gaining coverage.
#
# (lua_field, sentinel, scontrol_field, sbatch args, expected seeded value,
#  non-empty value the hook assigns to prove the setter is live)
#
# "{path}" in either the seed or the expected value is substituted with a
# per-test temporary path.  The assigned value cannot use it: the Lua hook
# is generated once at import, before any test has a working directory.
# ---------------------------------------------------------------------------
_GROUP_D_CASES = [
    # "" would revert these to a default, so each is seeded off it.
    (
        "partition",
        "TEST_PARTITION_CLEAR",
        "Partition",
        f"--partition={PARTITION_NAME}",
        PARTITION_NAME,
        "primary",
    ),
    (
        "qos",
        "TEST_QOS_CLEAR",
        "QOS",
        f"--account={ACCOUNT_NAME} --qos={QOS_NAME}",
        QOS_NAME,
        "normal",
    ),
    (
        "mail_user",
        "TEST_MAILUSER_CLEAR",
        "MailUser",
        "--mail-user=before@example.com --mail-type=END",
        "before@example.com",
        "after@example.com",
    ),
    # "" would clear these outright.
    (
        "std_out",
        "TEST_STDOUT_CLEAR",
        "StdOut",
        "--output={path}",
        "{path}",
        "test157_stdout.set",
    ),
    (
        "std_err",
        "TEST_STDERR_CLEAR",
        "StdErr",
        "--error={path}",
        "{path}",
        "test157_stderr.set",
    ),
    (
        "std_in",
        "TEST_STDIN_CLEAR",
        "StdIn",
        "--input={path}",
        "{path}",
        "test157_stdin.set",
    ),
]


# The same rows keyed by the sentinel that assigns a value instead of "".
_GROUP_D_SET_CASES = [
    (field, sentinel.replace("_CLEAR", "_SET"), show_field, seed, set_value)
    for field, sentinel, show_field, seed, _, set_value in _GROUP_D_CASES
]

# Each field gets a clearing branch and an assigning one, so the guardrail
# is always paired with proof that the setter it guards is reachable.
_GROUP_D_BRANCHES = [
    (field, sentinel, "") for field, sentinel, _, _, _, _ in _GROUP_D_CASES
] + [
    (field, sentinel, set_value)
    for field, sentinel, _, _, set_value in _GROUP_D_SET_CASES
]

_GROUP_D_LUA = "\n\n".join(
    f'    {"if" if index == 0 else "elseif"} c == "{sentinel}" then\n'
    f'        job_desc.{field} = "{value}"'
    for index, (field, sentinel, value) in enumerate(_GROUP_D_BRANCHES)
)

# ---------------------------------------------------------------------------
# Lua hook content - kept inline so the test is self-contained and
# require_config_file() can install it as the slurm-user.
# ---------------------------------------------------------------------------
_LUA_CONTENT = f"""--[[
 job_submit.lua for bug 25211 guardrail regression tests.

 Sentinel values (passed via --comment / scontrol update Comment=<sentinel>):

   Submit-path verification (Group C):
   TEST_SUBMIT_FEATURES_EMPTY    -- set features      = "" in slurm_job_submit
   TEST_SUBMIT_TRESPERNODE_EMPTY -- set tres_per_node = "" in slurm_job_submit
   TEST_SUBMIT_GRES_EMPTY        -- set gres          = "" in slurm_job_submit
   TEST_SUBMIT_COMMENT_EMPTY     -- set comment       = "" in slurm_job_submit

   The Group D branches are generated from _GROUP_D_CASES: per field, a
   _CLEAR sentinel assigning "" and a _SET sentinel assigning a value.

   Non-empty assignment (Group E):
   TEST_SET_FEATURES        -- set features      = "testfeat157"   (E1)
   TEST_SET_TRESPERNODE     -- set tres_per_node = "gres/gpu:1"    (E2)
   TEST_SET_COMMENT         -- set comment       = <new value>     (E3)
   TEST_SUBMIT_SET_FEATURES -- set features on the submit path     (E4)
   TEST_SUBMIT_SET_LAYOUT   -- set layout = "{LUA_LAYOUT}" on submit (E5)
   TEST_SUBMIT_READ_LAYOUT  -- log_user the submitted layout         (E5)
   TEST_SUBMIT_BAD_LAYOUT   -- set layout = "bad" on submit          (E5)
--]]

function slurm_job_submit(job_desc, part_list, submit_uid)
    -- Group C: on the submit path "" leaves the field NULL, so each of
    -- these assignments must discard the value the user submitted.
    if job_desc.comment == "TEST_SUBMIT_FEATURES_EMPTY" then
        job_desc.features = ""
    elseif job_desc.comment == "TEST_SUBMIT_TRESPERNODE_EMPTY" then
        job_desc.tres_per_node = ""
    elseif job_desc.comment == "TEST_SUBMIT_GRES_EMPTY" then
        job_desc.gres = ""
    elseif job_desc.comment == "TEST_SUBMIT_COMMENT_EMPTY" then
        job_desc.comment = ""

    elseif job_desc.comment == "TEST_SUBMIT_SET_FEATURES" then
        job_desc.features = "testfeat157"

    elseif job_desc.comment == "TEST_SUBMIT_SET_LAYOUT" then
        job_desc.layout = "{LUA_LAYOUT}"
    elseif job_desc.comment == "TEST_SUBMIT_READ_LAYOUT" then
        slurm.log_user("layout is %s", job_desc.layout or "nil")
    elseif job_desc.comment == "TEST_SUBMIT_BAD_LAYOUT" then
        job_desc.layout = "bad"
    end
    return slurm.SUCCESS
end

function slurm_job_modify(job_desc, job_rec, part_list, modify_uid)
    local c = job_desc.comment or ""

    -- Group D: these fields are not whitelisted, so "" drops to NULL inside
    -- the plugin, while a non-empty value still reaches the controller.
{_GROUP_D_LUA}

    -- Group E: a non-empty assignment must reach the controller.
    elseif c == "TEST_SET_FEATURES" then
        job_desc.features = "testfeat157"

    elseif c == "TEST_SET_TRESPERNODE" then
        job_desc.tres_per_node = "gres/gpu:1"

    elseif c == "TEST_SET_COMMENT" then
        job_desc.comment = "TEST_SET_COMMENT_DONE"

    end

    return slurm.SUCCESS
end

slurm.log_info("bug25211 guardrail test hook loaded")
return slurm.SUCCESS
"""


# ---------------------------------------------------------------------------
# Module-level fixture: install Lua plugin config and start Slurm.
# ---------------------------------------------------------------------------
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_config_parameter_includes("JobSubmitPlugins", "lua")
    atf.require_config_file("job_submit.lua", _LUA_CONTENT)
    # Accounting is required so a job's QOS is recorded and visible to
    # scontrol for the D2 guardrail.
    atf.require_accounting(modify=True)
    # A node feature and a fake GPU gres let the Group C jobs submit with
    # Features/TresPerNode already set, so the hook has a value to destroy.
    atf.require_tty(0)
    atf.require_config_parameter("Name", {"gpu": {"File": "/dev/tty0"}}, source="gres")
    atf.require_config_parameter_includes("GresTypes", "gpu")
    # An empty partition resolves to the cluster default, so the partition
    # job must sit in a non-default partition for that revert to be
    # observable.
    atf.require_config_parameter(
        "PartitionName",
        {
            "primary": {"Nodes": "ALL", "Default": "YES"},
            PARTITION_NAME: {"Nodes": "ALL"},
        },
    )
    atf.require_nodes(1, [("Gres", "gpu:1"), ("Features", "testfeat157")])
    atf.require_slurm_running()

    # The controller maps an empty qos to "revert to the user's default
    # QOS", so the qos job must hold a non-default QOS for that revert to
    # be observable.
    slurm_user = atf.properties["slurm-user"]
    test_user = atf.properties["test-user"]
    atf.run_command(f"sacctmgr -i add qos {QOS_NAME}", user=slurm_user, fatal=True)
    atf.run_command(
        f"sacctmgr -i add account {ACCOUNT_NAME}", user=slurm_user, fatal=True
    )
    atf.run_command(
        f"sacctmgr -i add user {test_user} account={ACCOUNT_NAME} "
        f"qos={QOS_NAME},normal",
        user=slurm_user,
        fatal=True,
    )

    yield

    atf.run_command(
        f"sacctmgr -i remove user {test_user} account={ACCOUNT_NAME}",
        user=slurm_user,
        quiet=True,
    )
    atf.run_command(
        f"sacctmgr -i remove account {ACCOUNT_NAME}", user=slurm_user, quiet=True
    )
    atf.run_command(f"sacctmgr -i remove qos {QOS_NAME}", user=slurm_user, quiet=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _submit_held(*, comment=None, extra_args=None):
    """Submit a held batch job and return its job id.

    extra_args is appended verbatim, so a case needing its own --output
    can pass one and it wins over the default.
    """
    args = "--hold -n1 --output=/dev/null --wrap='true'"
    if comment is not None:
        args += f" --comment={comment}"
    if extra_args:
        args += f" {extra_args}"
    return atf.submit_job_sbatch(args, fatal=True)


def _modify(job_id, **kwargs):
    """Run scontrol update on job_id with arbitrary key=value pairs."""
    parts = [f"{k}={v}" for k, v in kwargs.items()]
    return atf.run_command(
        f"scontrol update JobId={job_id} " + " ".join(parts),
        quiet=True,
    )


def _field(job_id, param):
    """Return get_job_parameter result, or None if absent or (null)."""
    return atf.get_job_parameter(job_id, param, default=None, quiet=True)


# ===========================================================================
# Group C - Submit-path behavior unchanged by the patch
# ===========================================================================
# job_submit_plugins.shtml documents that on the submit path assigning ""
# to a string field leaves it NULL, discarding whatever the user supplied,
# and that the modify path's per-field clear semantics do not apply there.


# The gres alias is sampled alongside tres_per_node because it is
# documented as clearable in its own right on the modify path.
#
# The comment case is the asymmetry at its sharpest: the same assignment
# that clears Comment on modify instead destroys the user's own --comment
# value here, which is the sentinel itself.
_GROUP_C_CASES = [
    ("features", "--constraint=testfeat157", "TEST_SUBMIT_FEATURES_EMPTY", "Features"),
    ("tres_per_node", "--gres=gpu:1", "TEST_SUBMIT_TRESPERNODE_EMPTY", "TresPerNode"),
    ("gres", "--gres=gpu:1", "TEST_SUBMIT_GRES_EMPTY", "TresPerNode"),
    ("comment", "", "TEST_SUBMIT_COMMENT_EMPTY", "Comment"),
]


@pytest.mark.parametrize(
    "seed, sentinel, show_field",
    [case[1:] for case in _GROUP_C_CASES],
    ids=[case[0] for case in _GROUP_C_CASES],
)
def test_C_submit_becomes_null(seed, sentinel, show_field):
    """On submit, assigning "" leaves the field NULL, dropping the user value."""
    job_id = atf.submit_job_sbatch(
        f"--hold -n1 --output=/dev/null {seed} --comment={sentinel} --wrap='true'",
        fatal=True,
    )
    after = _field(job_id, show_field)
    assert after is None, f"{show_field} must be null on submit path, got {after!r}"


# ===========================================================================
# Group D - Non-whitelist guardrails
# ===========================================================================


@pytest.mark.parametrize(
    "sentinel, show_field, seed, expected",
    [case[1:5] for case in _GROUP_D_CASES],
    ids=[case[0] for case in _GROUP_D_CASES],
)
def test_D_non_whitelist_empty_ignored(sentinel, show_field, seed, expected):
    """Lua assigns "" to a non-whitelisted field on modify -> field unchanged."""
    path = f"{os.getcwd()}/{show_field}.path"
    job_id = _submit_held(extra_args=seed.format(path=path))
    before = _field(job_id, show_field)
    want = expected.format(path=path)
    assert (
        before == want
    ), f"pre-condition: {show_field} must be {want!r}, got {before!r}"
    result = _modify(job_id, Comment=sentinel)
    assert result["exit_code"] == 0, (
        f'scontrol update must succeed when Lua sets {show_field} to "", '
        f"stderr: {result['stderr']}"
    )
    after = _field(job_id, show_field)
    assert (
        after == before
    ), f"{show_field} must be unchanged: expected {before!r}, got {after!r}"


# Without this, test_D_non_whitelist_empty_ignored would keep passing if the
# field's Lua setter were renamed or removed, since an assignment that never
# happens also leaves the field unchanged.
@pytest.mark.parametrize(
    "sentinel, show_field, seed, set_value",
    [case[1:] for case in _GROUP_D_SET_CASES],
    ids=[case[0] for case in _GROUP_D_SET_CASES],
)
def test_D_non_whitelist_set_works(sentinel, show_field, seed, set_value):
    """Lua assigns a non-empty value to a non-whitelisted field -> it applies."""
    path = f"{os.getcwd()}/{show_field}.path"
    job_id = _submit_held(extra_args=seed.format(path=path))
    result = _modify(job_id, Comment=sentinel)
    assert result["exit_code"] == 0, (
        f"scontrol update must succeed when Lua sets {show_field} to "
        f"{set_value!r}, stderr: {result['stderr']}"
    )
    # scontrol resolves a relative stdio path against the job's working
    # directory, which is where the job was submitted from.
    want = f"{os.getcwd()}/{set_value}" if "{path}" in seed else set_value
    after = _field(job_id, show_field)
    assert after == want, f"{show_field} must be {want!r}, got {after!r}"


# ===========================================================================
# Group E - Non-empty assignment still takes effect
# ===========================================================================
# job_submit_plugins.shtml documents job_desc as input/output and the hook
# as able to modify the job parameters supplied by the user.  Narrowing
# which fields "" clears must leave ordinary assignment working.

_GROUP_E_CASES = [
    ("features", "TEST_SET_FEATURES", "Features", "testfeat157"),
    ("tres_per_node", "TEST_SET_TRESPERNODE", "TresPerNode", "gres/gpu:1"),
    ("comment", "TEST_SET_COMMENT", "Comment", "TEST_SET_COMMENT_DONE"),
]


@pytest.mark.parametrize(
    "sentinel, show_field, expected",
    [case[1:] for case in _GROUP_E_CASES],
    ids=[case[0] for case in _GROUP_E_CASES],
)
def test_E_set_on_modify(sentinel, show_field, expected):
    """Lua assigns a non-empty value on modify -> the field takes that value."""
    job_id = _submit_held()
    before = _field(job_id, show_field)
    assert (
        before != expected
    ), f"pre-condition: {show_field} must not already be {expected!r}"
    result = _modify(job_id, Comment=sentinel)
    assert result["exit_code"] == 0, f"scontrol update failed: {result['stderr']}"
    after = _field(job_id, show_field)
    assert after == expected, f"{show_field} must be {expected!r}, got {after!r}"


def test_E_set_on_submit():
    """Lua assigns a non-empty value on submit -> the field takes that value."""
    job_id = atf.submit_job_sbatch(
        "--hold -n1 --output=/dev/null --comment=TEST_SUBMIT_SET_FEATURES "
        "--wrap='true'",
        fatal=True,
    )
    features_after = _field(job_id, "Features")
    assert (
        features_after == "testfeat157"
    ), f"Features must be set on submit path, got {features_after!r}"


@requires_layout
def test_E_set_layout_on_submit():
    """Lua assigns a layout on submit -> the job takes that layout."""
    job_id = _submit_held(comment="TEST_SUBMIT_SET_LAYOUT")
    layout = _field(job_id, "Layout")
    assert layout == LUA_LAYOUT, f"Layout must be {LUA_LAYOUT!r}, got {layout!r}"


@requires_layout
def test_E_read_layout_on_submit():
    """Lua reads the layout the user submitted."""
    result = atf.run_command(
        f"sbatch --hold -n1 --output=/dev/null --layout={USER_LAYOUT} "
        "--comment=TEST_SUBMIT_READ_LAYOUT --wrap='true'",
        fatal=True,
    )
    assert (
        f"layout is {USER_LAYOUT}" in result["stderr"]
    ), f"Lua should see the submitted layout, got: {result['stderr']}"


@requires_layout
def test_E_bad_layout_on_submit():
    """Lua assigns a malformed layout on submit -> the submit is refused."""
    # No fatal=True: the refusal is what is being asserted.
    result = atf.run_command(
        "sbatch --hold -n1 --output=/dev/null "
        "--comment=TEST_SUBMIT_BAD_LAYOUT --wrap='true'"
    )
    assert result["exit_code"] != 0, (
        "slurmctld must refuse the malformed layout Lua set, got stdout: "
        f"{result['stdout']}"
    )
    assert (
        "Invalid layout" in result["stderr"]
    ), f"The refusal should say the layout is invalid, got: {result['stderr']}"