| ############################################################################ | |
| # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| ############################################################################ | |
| import pytest | |
| import atf | |
| line_received = "Line received:" | |
| def setup_module(): | |
| atf.require_version( | |
| (26, 5), | |
| "bin/sattach", | |
| reason="Ticket 22180: SLUID availability added in 26.05+", | |
| ) | |
| atf.require_nodes(1) | |
| atf.require_slurm_running() | |
| @pytest.fixture(scope="function") | |
| def interactive_step_id(request): | |
| """ | |
| Submits a job with a single step waiting for 'trigger' input, and returns | |
| the StepId based on the request.param, so it can be JobId.0 or SLUID.0 | |
| """ | |
| file_out = "file.out" | |
| ready_str = "ready to read input" | |
| step_script = "interactive_step.sh" | |
| atf.make_bash_script( | |
| step_script, | |
| f""" | |
| echo {ready_str} | |
| read line | |
| echo "{line_received} $line" | |
| """, | |
| ) | |
| job_script = "job_with_interactive_step.sh" | |
| atf.make_bash_script( | |
| job_script, | |
| # We need to keep the stdin of the step open or an EOF will close it | |
| # and sattach won't be able to use it. | |
| f"""srun {step_script} < <(sleep infinity)""", | |
| ) | |
| job_id = atf.submit_job_sbatch(f"-t1 -o {file_out} {job_script}", fatal=True) | |
| # Wait until the interactive step is ready to read the line | |
| atf.wait_for_file(file_out) | |
| for t in atf.timer(): | |
| output = atf.run_command_output(f"cat {file_out}") | |
| if ready_str in output: | |
| break | |
| else: | |
| pytest.fail(f"File {file_out} wasn't ready") | |
| # Return the StepID using either the JobId or SLUID based on the param | |
| return f"{atf.get_job_parameter(job_id, request.param)}.0" | |
| @pytest.mark.parametrize("interactive_step_id", ["JobId", "SLUID"], indirect=True) | |
| def test_sattach(interactive_step_id): | |
| """Verify sattach receives output and can send input to the step, using JobId or SLUID.""" | |
| result = atf.run_command(f"sattach {interactive_step_id}", input="trigger") | |
| assert ( | |
| f"{line_received} trigger" in result["stdout"] | |
| ), f"sattach should attach IO streams and get '{line_received} trigger', got: {result['stdout']}" | |
| assert result["exit_code"] == 0, "sattach should end normally" |