############################################################################
# Copyright (C) SchedMD LLC.
############################################################################
import atf
import os
import pytest

# import re

iteration_count = 100
node_count = 4
task_count = 4


# Setup/Teardown
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_nodes(node_count, [("CPUs", task_count)])
    atf.require_slurm_running()


@pytest.fixture(scope="function")
def input_file(request, tmp_path):
    """Create a sizable text file"""
    in_file = str(tmp_path / "test.in")
    with open(in_file, "w") as f:
        for line in range(0, 100):
            for i in range(1, 100):
                f.write(chr(ord("0") + i % 10))
            f.write("\n")
    return in_file


def test_stdin_broadcast(input_file, tmp_path):
    """Copy job input to job output and compare sizes"""
    output_file = str(tmp_path / "test.out")
    for i in range(0, iteration_count):
        if os.path.isfile(output_file):
            os.remove(output_file)
        exit_code = atf.run_command_exit(
            f"srun --input={input_file} --output={output_file} --error=/dev/null -n {task_count} -N {node_count} -t 1 cat -"
        )
        assert exit_code == 0, f"srun failed with status {exit_code}"
        in_size = os.path.getsize(input_file)
        out_size = os.path.getsize(output_file)
        assert (
            out_size == in_size * task_count
        ), f"output file size ({out_size}) was not {task_count} times the input file size ({in_size})"