unsloth/tests/python/test_docker_studio_launch_view_config.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0

"""A quote in UNSLOTH_NOTEBOOKS_VIEW_DIR made the generated jupyter config a SyntaxError."""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import sys
import types
from pathlib import Path

import pytest
import urllib.parse

REPO_ROOT = Path(__file__).resolve().parents[2]
LAUNCH = REPO_ROOT / "docker" / "studio_launch.sh"

behavioural = pytest.mark.skipif(shutil.which("bash") is None, reason = "needs bash")


@pytest.fixture(scope = "module")
def generator() -> str:
    """The `python - >> ... <<'PY' ... PY` block, verbatim from the script."""
    source = LAUNCH.read_text(encoding = "utf-8")
    match = re.search(
        r"^\s*UNSLOTH_VIEW_REL=.*?\n\s*python - >> .*?<<'PY'\n(.*?)\nPY$",
        source,
        re.S | re.M,
    )
    assert match, "the notebook-view config generator disappeared or changed shape"
    return match.group(1)


def _render(generator: str, rel: str, view: str) -> str:
    result = subprocess.run(
        [sys.executable, "-"],
        input = generator,
        capture_output = True,
        text = True,
        timeout = 120,
        env = dict(os.environ, UNSLOTH_VIEW_REL = rel, UNSLOTH_VIEW_DIR = view),
    )
    assert result.returncode == 0, result.stdout + result.stderr
    return result.stdout


def _load(config_text: str):
    """Execute the config the way Jupyter does, and hand back what it set."""

    class _Node(types.SimpleNamespace):
        pass

    config = types.SimpleNamespace(
        ServerApp = _Node(), LabApp = _Node(), PasswordIdentityProvider = _Node()
    )
    exec(compile(config_text, "jupyter_lab_config.py", "exec"), {"c": config})
    return config


NASTY = [
    pytest.param('My "Special" Notebooks', id = "double-quote"),
    pytest.param("Note\\tbooks", id = "backslash-t"),
    pytest.param("Note'books", id = "single-quote"),
    pytest.param("Note\\books", id = "trailing-backslash-segment"),
    pytest.param("Unsloth Notebooks", id = "the-default"),
]


@pytest.mark.parametrize("name", NASTY)
def test_the_config_stays_valid_python_and_keeps_the_path(generator: str, name: str):
    view = f"/workspace/{name}"
    rendered = _render(generator, name, view)

    config = _load(rendered)

    assert (
        config.ServerApp.preferred_dir == view
    ), "the path Jupyter ends up with must be the one the user asked for"
    # default_url must be URL-quoted: the default view directory contains a space.
    assert config.ServerApp.default_url == "/lab/tree/" + urllib.parse.quote(name)
    assert " " not in config.ServerApp.default_url, "an unencoded space breaks the printed link"
    assert (
        config.LabApp.default_url == config.ServerApp.default_url
    ), "LabApp otherwise overrides ServerApp back to /lab"


def test_a_newline_in_the_path_cannot_inject_a_config_line(generator: str):
    name = 'x"\nc.ServerApp.token = "pwned'
    config = _load(_render(generator, name, f"/workspace/{name}"))

    assert config.ServerApp.preferred_dir == f"/workspace/{name}"
    assert not hasattr(config.ServerApp, "token"), "the value was executed as config"


def test_the_generator_does_not_interpolate_the_paths_in_the_shell(generator: str):
    # Heredoc must stay quoted with values via env, or shell expansion undoes python quoting.
    source = LAUNCH.read_text(encoding = "utf-8")
    assert "python - >> \"${JUPYTER_CONFIG_DIR}/jupyter_lab_config.py\" <<'PY'" in source
    assert "${_view_rel}" not in generator and "${_view_dir}" not in generator
    assert 'os.environ["UNSLOTH_VIEW_REL"]' in generator
    assert 'os.environ["UNSLOTH_VIEW_DIR"]' in generator


@behavioural
def test_the_heredoc_form_this_replaced_really_was_broken(tmp_path: Path):
    """Pin the premise, so the test above is not guarding a hypothetical.

    Reproduces the previous construct exactly: an unquoted heredoc interpolating
    the value into a Python string literal. With a double quote in the path the
    generated config does not parse, which is JupyterLab failing to start.
    """
    name = 'My "Special" Notebooks'
    script = tmp_path / "old.sh"
    script.write_text(
        "#!/usr/bin/env bash\n"
        '_view_rel="$1"\n'
        '_view_dir="$2"\n'
        "cat <<EOF\n"
        'c.ServerApp.default_url = "/lab/tree/${_view_rel}"\n'
        'c.LabApp.default_url = "/lab/tree/${_view_rel}"\n'
        'c.ServerApp.preferred_dir = "${_view_dir}"\n'
        "EOF\n",
        encoding = "utf-8",
    )
    result = subprocess.run(
        ["bash", str(script), name, f"/workspace/{name}"],
        capture_output = True,
        text = True,
        timeout = 120,
    )
    assert result.returncode == 0, result.stderr

    with pytest.raises(SyntaxError):
        compile(result.stdout, "jupyter_lab_config.py", "exec")