spawner/tests/test_spawner.py

"""Test GLiCID spawner module."""

import re
from collections import namedtuple
from itertools import repeat

import glicid_spawner.spawner
import pytest
from batchspawner import SlurmSpawner, format_template
from glicid_spawner import GlicidSpawner
from glicid_spawner.spawner import asyncio

User = namedtuple('User', 'name,url')
Hub = namedtuple('Hub', 'public_host,api_url,base_url')


def test_spawner_config():
    """Test spawner configuration."""
    spawner = GlicidSpawner(progress_rate=10)

    assert isinstance(spawner, GlicidSpawner)
    assert isinstance(spawner, SlurmSpawner)

    assert spawner.batchspawner_singleuser_cmd == 'glicid-spawner-singleuser'

    cmd = spawner.cmd_formatted_for_batch()
    assert 'bin/glicid-spawner-singleuser' in cmd
    assert 'bin/jupyterhub-singleuser' in cmd

    assert spawner.req_mamba_exe == '/micromamba/operator/bin/micromamba'
    assert spawner.req_mamba_root_prefix == '/micromamba/operator'
    assert spawner.req_mamba_user_base == '/micromamba/$USER'
    assert spawner.req_job_name == 'jupyterhub_glicid'
    assert spawner.req_qos == 'short'

    assert spawner.disable_user_config
    assert spawner.notebook_dir == '/'
    assert spawner.default_url == '/lab/tree/home/{username}?reset'
    assert spawner.progress_rate == 10


def test_spawner_batch_script(monkeypatch):
    """Test spawner SLURM batch script."""
    monkeypatch.setattr(GlicidSpawner, 'get_env', lambda _: {})

    spawner = GlicidSpawner(
        user=User('john-doe', '/john-doe/'),
        req_homedir='/home/john-doe',
        hub=Hub('public.hostname', 'api.url', 'base.url'),
        user_options={
            'cluster': 'nautilus',
            'workdir': '/home/john-doe',
            'pyenv': '/micromamba/john-doe/envs/foo',
        },
    )

    script = format_template(
        spawner.batch_script,
        cmd=spawner.cmd_formatted_for_batch(),
        **(spawner.get_req_subvars() | spawner.user_options),
    )

    assert '#SBATCH --job-name=jupyterhub_glicid' in script
    assert '#SBATCH --chdir=/home/john-doe' in script
    assert '#SBATCH --cluster=nautilus' in script
    assert '#SBATCH --qos=short' in script

    assert 'scontrol write batch_script ${SLURM_JOB_ID} --cluster=nautilus -;' in script

    assert 'export JUPYTER_LOG_DIR="/home/john-doe/.jupyter/spawner/logs"' in script

    assert 'export MAMBA_EXE=/micromamba/operator/bin/micromamba;' in script
    assert 'export MAMBA_ROOT_PREFIX=/micromamba/operator;' in script
    assert 'export PYTHONUSERBASE=/micromamba/$USER;' in script
    assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script

    assert 'micromamba activate /micromamba/john-doe/envs/foo;' in script
    assert 'export JUPYTER_PATH=/micromamba/john-doe/envs/foo/share/jupyter;' in script

    assert re.search(r'.*/bin/glicid-spawner-singleuser .*/bin/jupyterhub-singleuser', script)


def test_spawner_parse_job_id():
    """Test spawner job id parser."""
    spawner = GlicidSpawner()

    assert spawner.parse_job_id('123') == '123'
    assert spawner.parse_job_id('456;nautilus') == '456 -M nautilus'

    assert spawner.parse_job_id('') == ''


def test_spawner_options_form(monkeypatch):
    """Test spawner options form."""
    monkeypatch.setattr(
        glicid_spawner.spawner, 'options_form', lambda username: f"options_form('{username}')"
    )
    monkeypatch.setattr(
        glicid_spawner.spawner,
        'options_from_form',
        lambda formdata: f'options_from_form({formdata})',
    )

    spawner = GlicidSpawner(user=User('john-doe', '/john-doe/'))

    assert spawner.user.name == 'john-doe'
    assert spawner.options_form == "options_form('john-doe')"
    assert spawner.options_from_form({'foo': 123}) == "options_from_form({'foo': 123})"

    spawner = GlicidSpawner(user_options={'workdir': '/LAB-DATA'})

    assert spawner.default_url == '/lab/tree/LAB-DATA?reset'


@pytest.mark.asyncio
async def test_spawner_progress(monkeypatch):
    """Test spawner progress messages."""

    def isrunning():
        """Running generator values."""
        yield False  # NOTFOUND
        yield False  # PENDING
        yield from repeat(True)  # RUNNING

    def ispending():
        """Pending generator values."""
        yield False  # NOTFOUND
        yield True  # PENDING
        yield from repeat(False)  # RUNNING (never used)

    # Mock Job status generator
    iter_ispending = iter(ispending())
    iter_isrunning = iter(isrunning())

    monkeypatch.setattr(
        glicid_spawner.GlicidSpawner, 'state_isrunning', lambda _: next(iter_isrunning)
    )
    monkeypatch.setattr(
        glicid_spawner.GlicidSpawner, 'state_ispending', lambda _: next(iter_ispending)
    )

    async def mock_sleep(_):
        """Mock asyncio sleep."""

    monkeypatch.setattr(asyncio, 'sleep', mock_sleep)

    spawner = GlicidSpawner(progress_rate=20)

    progress = [msg['progress'] async for msg in spawner.progress()]

    assert progress == [
        10,  # submit
        20,  # pending
        60,  # running | elapse time = 20 -> setup
        80,  # running | elapse time = 40 -> connect
        95,  # running | elapse time = 60 -> too long
    ]