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

# import re

approved_account = "account1"
unapproved_account = "account2"
test_user = pwd.getpwuid(os.getuid())[0]


# Setup
@pytest.fixture(scope="module", autouse=True)
def setup():
    atf.require_accounting(modify=True)
    atf.require_config_parameter("AccountingStorageEnforce", "associations")
    atf.require_slurm_running()


@pytest.fixture(scope="module")
def setup_account():
    atf.run_command(
        f"sacctmgr -vi add account {approved_account}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )
    atf.run_command(
        f"sacctmgr -vi add user {test_user} account={approved_account}",
        user=atf.properties["slurm-user"],
        fatal=True,
    )


def test_approved_association(setup_account):
    """Verify that a job can be submitted under an approved account"""
    assert atf.run_command_exit(f"srun --account={approved_account} true") == 0


def test_unapproved_association():
    """Verify that a job cannot be submitted under an unapproved account"""
    assert atf.run_command_exit(f"srun --account={unapproved_account} true") != 0