#!/usr/bin/env pytest

# Test that we can create sub-cgroups.  If not, skip the test
# if so, turn on the secret unprivileged cgroup option,
# and submit a job that uses more memory than we have
# requested.  Verify that it goes on hold.

import logging
import pytest
import os

import htcondor2 as htcondor
from ornithology import *

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# Standup a condor with user-space cgroup v2 enabled
@standup
def condor(test_dir):
    with Condor(test_dir / "condor", config={
        "CREATE_CGROUP_WITHOUT_ROOT": "true",
        "LOCAL_UNIVERSE_CGROUP_ENFORCEMENT": "true",
        }) as condor:
        yield condor

# Standup a condor with user-space cgroup v2 enabled and job hook
@standup
def condor_with_hook(test_dir):
    with Condor(test_dir / "condor_with_hooks", config={"CREATE_CGROUP_WITHOUT_ROOT": "true",
                                                        "STARTER_DEFAULT_JOB_HOOK_KEYWORD": "TEST",
                                                        "TEST_HOOK_PREPARE_JOB": test_dir / "hook_starter_test_prepare.sh"}) as condor_with_hook:
        yield condor_with_hook

# Standup a condor with user-space cgroup v2 enabled and SUBMIT_REQUEST_MISSING_UNITS=error
@standup
def condor_with_units_error(test_dir):
    with Condor(test_dir / "condor_units_error", config={"CREATE_CGROUP_WITHOUT_ROOT": "true",
                                                          "SUBMIT_REQUEST_MISSING_UNITS": "error"}) as condor_units:
        yield condor_units

# The actual job will allocate 200 Mb in python, but request 50Mb.
# The 50Mb might get rounded up to 128, so be careful.
@action
def job_script_contents():
    return format_script( """
#!/usr/bin/env python3

import os
import time

two_hundred_meg = bytearray(200 * 1024 * 1024)
for i in range(10):
    time.sleep(1)
exit(0)
""")

@action
def test_script(test_dir, job_script_contents):
    test_script = test_dir / "job_script.py"
    write_file(test_script, job_script_contents)
    return test_script

@action
def test_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "oom_log",
            "keep_claim_idle": "300",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def completed_test_job(condor, test_job_hash):
    ctj = condor.submit(
        {**test_job_hash}, count=1
    )
    assert ctj.wait(
        condition=ClusterState.all_held,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_complete,
    )
    return ctj

@action
def escaping_job_hash(test_dir, path_to_python, test_script):
    # A job which spawns "sleep 1234", with all the condor tracking env vars removed
    # and nohup'd, to avoid all non-cgroup tracking mechanisms.  Writes the pid
    # to pid file, which lives in the test_dir.  If the pid in pidfile exists
    # after the job exits, the sleep escaped the job, and the test should fail.
    return {
            "shell": "nohup /usr/bin/env -i /usr/bin/sleep 1234 & echo $! > pidfile",
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "escaped_log",
            "keep_claim_idle": "300",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def escaping_test_job(condor, escaping_job_hash):
    ej = condor.submit(
        {**escaping_job_hash}, count=1
    )
    assert ej.wait(
        condition=ClusterState.all_complete,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_held,
    )
    return ej

@action
def isolated_job_hash(test_dir, path_to_python):
    # A job which cats the content of its own cgroup.procs.  Because we
    # exec, the shell has gone away, so the cgroup should only
    # contain the cat process, which verifies that the job is isolated into
    # its own cgroup.
    return {
            "shell": "exec /usr/bin/cat /sys/fs/cgroup$(awk -F: '/^0:/ {print $3}' < /proc/self/cgroup)/cgroup.procs > cgroup.procs",
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "isolated_log",
            "keep_claim_idle": "300",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }


@action
def isolated_test_job(condor, isolated_job_hash):
    ij = condor.submit(
        {**isolated_job_hash}, count=1
    )
    assert ij.wait(
        condition=ClusterState.all_complete,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_held,
    )
    return ij

@action
def slice_scope_job_hash(test_dir, path_to_python):
    # A job which cats the content of /proc/self/cgroup.
    # This can be interrogated to verify that the job is the right cgroup.  
    return {
            "shell": "/usr/bin/cat /proc/self/cgroup > cgroup",
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "slice_scope_log",
            "keep_claim_idle": "300",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def slice_scope_test_job(condor, slice_scope_job_hash):
    ssj = condor.submit(
        {**slice_scope_job_hash}, count=1
    )
    assert ssj.wait(
        condition=ClusterState.all_complete,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_held,
    )
    return ssj

@action
def memory_detection_job_hash(test_dir, condor, path_to_python):
    # A job which runs condor_config_val DETECTED_MEMORY.
    # This should return the memory limit of the cgroup
    bin_dir = condor.run_command( ['condor_config_val', 'BIN']).stdout
    return {
            "shell": f"{bin_dir}/condor_config_val DETECTED_MEMORY",
            "environment": "CONDOR_CONFIG=/dev/null",
            "universe": "vanilla",
            "output": "detected_memory",
            "error": "error",
            "log": "detected_memory.log",
            "keep_claim_idle": "300",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def hook_script_contents():
    return format_script( """
#!/bin/bash
echo 'HookStatusCode = 0'
exit 0
""")

@action
def hook_script(test_dir, hook_script_contents):
    hook_script = test_dir / "hook_starter_test_prepare.sh"
    write_file(hook_script, hook_script_contents)
    os.chmod(hook_script, 0o755)
    return hook_script

@action
def hook_test_job_hash(test_dir, hook_script, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "hook_log",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def memory_detection_test_job(condor, memory_detection_job_hash):
    mmj = condor.submit(
        {**memory_detection_job_hash}, count=1
    )
    assert mmj.wait(
        condition=ClusterState.all_complete,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_held,
    )
    return mmj

@action
def hook_test_job(condor_with_hook, hook_test_job_hash):
    htj = condor_with_hook.submit(
        {**hook_test_job_hash}, count=1
    )
    assert htj.wait(
        condition=ClusterState.all_held,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_complete,
    )
    return htj

@action
def retry_memory_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "retry_log",
            "keep_claim_idle": "300",
            "request_memory": "50mb",
            "retry_request_memory": "400mb",
            "request_cpus": "1",
            "request_disk": "200mb"
            }

@action
def retry_memory_job(condor, retry_memory_job_hash):
    ctj = condor.submit(
        {**retry_memory_job_hash}, count=1
    )
    assert ctj.wait(
        condition=ClusterState.all_complete,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_held,
    )
    return ctj

@action
def retry_memory_no_units_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "retry_no_units_log",
            "keep_claim_idle": "300",
            "request_memory": "50mb",
            "retry_request_memory": "400",  # No units - should fail
            "request_cpus": "1",
            "request_disk": "200mb"
            }

@action
def retry_memory_max_no_units_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "retry_max_no_units_log",
            "keep_claim_idle": "300",
            "request_memory": "50mb",
            "retry_request_memory_max": "500",  # No units - should fail
            "request_cpus": "1",
            "request_disk": "200mb"
            }

@action
def retry_memory_with_units_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "retry_with_units_log",
            "keep_claim_idle": "300",
            "request_memory": "50mb",
            "retry_request_memory": "400mb",  # With units - should succeed
            "request_cpus": "1",
            "request_disk": "200mb"
            }

@action
def retry_memory_max_with_units_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "vanilla",
            "output": "output",
            "error": "error",
            "log": "retry_max_with_units_log",
            "keep_claim_idle": "300",
            "request_memory": "50mb",
            "retry_request_memory_max": "500mb",  # With units - should succeed
            "request_cpus": "1",
            "request_disk": "200mb"
            }

@action
def local_universe_job_hash(test_dir, path_to_python, test_script):
    return {
            "executable": path_to_python,
            "arguments": test_script,
            "universe": "local",
            "output": "local_output",
            "error": "local_error",
            "log": "local_oom_log",
            "request_memory": "50m",
            "request_cpus": "1",
            "request_disk": "200m"
            }

@action
def completed_local_test_job(condor, local_universe_job_hash):
    clj = condor.submit(
        {**local_universe_job_hash}, count=1
    )
    assert clj.wait(
        condition=ClusterState.all_held,
        timeout=120,
        verbose=True,
        fail_condition=ClusterState.any_complete,
    )
    return clj

# These tests only works if we start in a writeable (delegated)
# cgroup.  Return true if this is so
def CgroupIsWriteable():
    if os.environ.get("HTCONDOR_TEST_USE_SUDO") == "1":
        # When using sudo, we should have root
        return True

    with open('/proc/self/cgroup') as fd:
        proc_self_cgroup = fd.read()

    # proc_self_cgroup on a cgroup v2 system should contain
    # ::/path/to/cgroup/without/leading/sys/fs/cgroup
    my_cgroup = proc_self_cgroup.split(':')[2].rstrip()
    cgroup_filepath = "/sys/fs/cgroup" + my_cgroup
    cgroup_filepath_parent = cgroup_filepath + "/.."

    # technically, also need to verify that parent's cgroup.subtree_control
    # contains memory, pid, cpu, but let's skip that for now
    return os.access(cgroup_filepath_parent, os.W_OK)


@pytest.mark.skipif(not CgroupIsWriteable(), reason="Test was not born in a writeable cgroup v2")
class TestCgroupOOM:
    def test_cgroup_oom(self, completed_test_job):
        assert completed_test_job.state.all_held()

    def test_cgroup_escape(self, escaping_test_job):
        # When the job exits, pidfile contains the pid of the sleep job
        with open('pidfile', 'r') as file:
            pid = file.read().rstrip()
        # Check that the pid does not exist in the system anymore -- that it has not escaped.
        assert not os.path.exists("/proc/" + pid)

    def test_cgroup_isolated(self, isolated_test_job):
        # When the job exits, this is a copy of the cgroup.procs file
        # from the cgroup of the job
        cgroup_procs = "cgroup.procs"
        assert os.path.exists(cgroup_procs)
        with open(cgroup_procs, 'r') as file:
            for count, pid in enumerate(file):
                pass
        count += 1 # Python has 0-indexing
        assert count == 1

    def test_cgroup_slice_scope(self, slice_scope_test_job):
        # When the job exits, cgroup contains the cgroup of the cat job
        with open('cgroup', 'r') as file:
            line        = file.read().rstrip()
            ultimate    = line.split('/')[-1]
            penultimate = line.split('/')[-2]
            assert ultimate.endswith(".scope")
            assert penultimate.endswith(".slice")
        assert True

    def test_cgroup_memory_detection(self, memory_detection_test_job):
        # When the job exits, detected_memory contains condor's idea of memory
        memory = 0
        with open('detected_memory', 'r') as file:
            memory = int(file.read().rstrip())
        assert(50 < memory and memory < 200) # Allow for rounding up to 128 Mb

    def test_cgroup_with_hook(self, hook_test_job):
        assert hook_test_job.state.all_held()

    def test_cgroup_with_retry_memory(self, retry_memory_job):
        assert retry_memory_job.state.all_complete()

    def test_retry_memory_without_units_fails(self, condor_with_units_error, retry_memory_no_units_job_hash):
        # Test that submitting retry_request_memory without units fails when SUBMIT_REQUEST_MISSING_UNITS=error
        with pytest.raises(Exception):
            condor_with_units_error.submit({**retry_memory_no_units_job_hash}, count=1)

    def test_retry_memory_max_without_units_fails(self, condor_with_units_error, retry_memory_max_no_units_job_hash):
        # Test that submitting retry_request_memory_max without units fails when SUBMIT_REQUEST_MISSING_UNITS=error
        with pytest.raises(Exception):
            condor_with_units_error.submit({**retry_memory_max_no_units_job_hash}, count=1)

    def test_retry_memory_with_units_succeeds(self, condor_with_units_error, retry_memory_with_units_job_hash):
        # Test that submitting retry_request_memory with units succeeds even when SUBMIT_REQUEST_MISSING_UNITS=error
        job = condor_with_units_error.submit({**retry_memory_with_units_job_hash}, count=1)
        assert job.wait(
            condition=ClusterState.all_complete,
            timeout=120,
            verbose=True,
            fail_condition=ClusterState.any_held,
        )

    def test_local_universe_oom(self, completed_local_test_job):
        # Test that a local universe job that exceeds memory is put on hold
        assert completed_local_test_job.state.all_held()
