#!/usr/bin/env python3 # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 """Cross-platform validation of the Unsloth Docker JupyterLab/notebook features. Runs WITHOUT Docker or a GPU, so it can execute on the Linux/macOS/Windows CI lanes. It exercises the actual notebook-helper logic (not just py_compile) and checks the shipped JupyterLab config + labextension source, so a regression in the notebook organisation, Colab compatibility, Colab-intro/widget stripping, sidecar-log gating, the labextension plugins, the JupyterLab defaults, or the login branding fails CI on every device. Usage: python tests/validate_studio_features.py Exit 0 = all checks pass; non-zero = at least one failed. """ from __future__ import annotations import importlib import json import os import sys ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) DOCKER = os.path.join(ROOT, "docker") JUPYTER = os.path.join(DOCKER, "jupyter") LABEXT = os.path.join(JUPYTER, "unsloth_labext") sys.path.insert(0, DOCKER) _failures: list[str] = [] def check( name: str, cond: bool, detail: str = "", ) -> None: status = "PASS" if cond else "FAIL" print(f" [{status}] {name}" + (f" -- {detail}" if detail and not cond else "")) if not cond: _failures.append(name) # 1. Colab cell-magic compatibility (#@title then %%capture) def test_colab_compat() -> None: print("colab cell-magic compat (unsloth_colab_compat):") m = importlib.import_module("unsloth_colab_compat") out = m.colab_cell_magic_fix(["#@title Setup\n", "%%capture\n", "!pip install x\n"]) check("magic hoisted above #@title", out[0] == "%%capture\n" and "#@title Setup\n" in out) # idempotent / already on top same = ["%%capture\n", "print(1)\n"] check("no-op when magic already first", m.colab_cell_magic_fix(same) == same) # non-magic cell untouched plain = ["x = 1\n", "y = 2\n"] check("plain cell untouched", m.colab_cell_magic_fix(plain) == plain) # content magic (%%writefile) NOT hoisted into the written file body wf = ["#@title Config\n", "%%writefile config.json\n", "{}\n"] check("content magic (%%writefile) left untouched", m.colab_cell_magic_fix(wf) == wf) # safe magic with arg still hoisted bash = ["#@title Run\n", "%%bash\n", "echo hi\n"] check("safe magic (%%bash) hoisted", m.colab_cell_magic_fix(bash)[0] == "%%bash\n") # 2. Notebook categorisation (clean_section) + README parsing def test_nb_view() -> None: print("notebook view (unsloth_nb_view):") v = importlib.import_module("unsloth_nb_view") check( "clean_section dash/slash -> space", v.clean_section("### GRPO-Reinforcement/Learning Notebooks") == "GRPO Reinforcement Learning Notebooks", v.clean_section("### GRPO-Reinforcement/Learning Notebooks"), ) check( "clean_section strips hashes/space", v.clean_section("## Main Notebooks ") == "Main Notebooks", ) # 3. Colab-intro + stale-widget stripping def test_strip() -> None: print("notebook strip (unsloth_nb_strip_colab):") s = importlib.import_module("unsloth_nb_strip_colab") nb = { "metadata": {"widgets": {"application/vnd.jupyter.widget-state+json": {"x": 1}}}, "cells": [ { "cell_type": "markdown", "source": [ 'To run this, press "Runtime" ... Tesla T4 Google Colab instance!\n', "\n", "You will learn how to ...\n", ], }, { "cell_type": "code", "source": ["print(1)\n"], "outputs": [ {"output_type": "stream", "name": "stdout", "text": "ok\n"}, { "output_type": "display_data", "data": { "application/vnd.jupyter.widget-view+json": {"model_id": "abc"}, "text/plain": "0%| | 0/10", }, }, ], }, ], } changed1 = s._strip_intro(nb) changed2 = s._clean_widgets(nb) check( "intro line stripped", changed1 and not any("to run this, press" in (l.lower()) for l in nb["cells"][0]["source"]), ) check("intro body kept", any("You will learn" in l for l in nb["cells"][0]["source"])) wv = sum( 1 for c in nb["cells"] for o in (c.get("outputs", []) or []) if "application/vnd.jupyter.widget-view+json" in (o.get("data", {}) or {}) ) check("widget-view outputs removed", changed2 and wv == 0) check( "non-widget outputs kept", any( o.get("output_type") == "stream" for c in nb["cells"] for o in (c.get("outputs", []) or []) ), ) check("metadata.widgets removed", "widgets" not in nb["metadata"]) # idempotent check("strip idempotent", not s._strip_intro(nb) and not s._clean_widgets(nb)) # 4. Sidecar-log gating def test_sidecar_log_gate() -> None: print("sidecar log gate (unsloth_nb_compat):") c = importlib.import_module("unsloth_nb_compat") old = os.environ.pop("UNSLOTH_ENABLE_LOGGING", None) try: check("logging off by default", c._logging_enabled() is False) os.environ["UNSLOTH_ENABLE_LOGGING"] = "1" check("logging on with env=1", c._logging_enabled() is True) os.environ["UNSLOTH_ENABLE_LOGGING"] = "0" check("logging off with env=0", c._logging_enabled() is False) finally: os.environ.pop("UNSLOTH_ENABLE_LOGGING", None) if old is not None: os.environ["UNSLOTH_ENABLE_LOGGING"] = old # 5. JupyterLab defaults (overrides.json) def test_overrides() -> None: print("jupyterlab defaults (jupyter/overrides.json):") path = os.path.join(JUPYTER, "overrides.json") check("overrides.json exists", os.path.isfile(path)) if not os.path.isfile(path): return with open(path, encoding = "utf-8") as f: d = json.load(f) # raises -> CI fails if invalid JSON themes = d.get("@jupyterlab/apputils-extension:themes", {}) check( "default theme = Unsloth Dark", themes.get("theme") == "Unsloth Dark", str(themes.get("theme")), ) check("adaptive theme on", themes.get("adaptive-theme") is True) check("preferred dark = Unsloth Dark", themes.get("preferred-dark-theme") == "Unsloth Dark") tracker = d.get("@jupyterlab/notebook-extension:tracker", {}) check( "windowingMode none", tracker.get("windowingMode") == "none", str(tracker.get("windowingMode")), ) notif = d.get("@jupyterlab/apputils-extension:notification", {}) check( "news prompt off", str(notif.get("fetchNews")) == "false" and notif.get("checkForUpdates") is False, ) panel = d.get("@jupyterlab/notebook-extension:panel", {}) labels = [t.get("label", "") for t in panel.get("toolbar", [])] check( "Restart & Run All label (single >>)", any(l == "Restart & Run All" for l in labels) and not any(">>" in l for l in labels), str(labels), ) # 6. Labextension source (plugins) + login branding assets def test_labext_and_branding() -> None: print("labextension + branding assets:") pkg = os.path.join(LABEXT, "package.json") check("labext package.json exists", os.path.isfile(pkg)) if os.path.isfile(pkg): with open(pkg, encoding = "utf-8") as f: p = json.load(f) check("labext name unsloth-jupyterlab", p.get("name") == "unsloth-jupyterlab") check("labext themePath set", bool(p.get("jupyterlab", {}).get("themePath"))) # Concatenate every .ts module under src/ so plugins defined in their own # files (cellNav, colabTitle, outputSelect, uiChrome) are all covered. src_dir = os.path.join(LABEXT, "src") all_src = "" if os.path.isdir(src_dir): for fn in sorted(os.listdir(src_dir)): if fn.endswith(".ts"): with open(os.path.join(src_dir, fn), encoding = "utf-8") as f: all_src += f.read() + "\n" for plug in [ "unsloth-jupyterlab:theme", "unsloth-jupyterlab:cell-nav", "unsloth-jupyterlab:logo", "unsloth-jupyterlab:colab-title", "unsloth-jupyterlab:output-select-all", "unsloth-jupyterlab:ui-chrome", ]: check(f"plugin present: {plug}", plug in all_src) # The two newest plugins are also exported from index.ts (wired in). index = os.path.join(src_dir, "index.ts") index_src = open(index, encoding = "utf-8").read() if os.path.isfile(index) else "" check("outputSelect wired in index.ts", "outputSelectPlugin" in index_src) check("uiChrome wired in index.ts", "uiChromePlugin" in index_src) # uiChrome hides the right activity bar; CTRL+A output-select selects nodes. check("right activity bar hidden", "jp-mod-right" in all_src and "display: none" in all_src) check("ctrl+A output select", "selectNodeContents" in all_src) # The remembered pointer-down is only replaced by another pointer-down, but # J/K/arrow cell navigation fires none, so it has to be revalidated (still in # the document, still in the ACTIVE cell) before it is used as the fallback -- # otherwise Ctrl+A on a later cell selects the old output and swallows # JupyterLab's notebook:select-all. check( "ctrl+A fallback revalidated", "isConnected" in all_src and "jp-mod-active" in all_src, ) # branding assets login = os.path.join(JUPYTER, "login.html") login_src = open(login, encoding = "utf-8").read() if os.path.isfile(login) else "" check("login.html branded", "unsloth-login-card" in login_src) check( "login.html uses sloth stickers", 'static_url("sloth/' in login_src or "static_url('sloth/" in login_src, ) check("favicon.ico present", os.path.isfile(os.path.join(JUPYTER, "favicon.ico"))) check("logo.png present", os.path.isfile(os.path.join(JUPYTER, "logo.png"))) check( "sloth sticker installer present", os.path.isfile(os.path.join(JUPYTER, "install_sloth_stickers.py")), ) def main() -> int: print("=== Unsloth Studio/notebook feature validation ===") for t in ( test_colab_compat, test_nb_view, test_strip, test_sidecar_log_gate, test_overrides, test_labext_and_branding, ): try: t() except Exception as e: # a thrown exception is a failure, not a crash _failures.append(f"{t.__name__}: {e!r}") print(f" [FAIL] {t.__name__} raised {e!r}") print() if _failures: print(f"FAILED ({len(_failures)}): " + ", ".join(_failures)) return 1 print("ALL CHECKS PASSED") return 0 if __name__ == "__main__": sys.exit(main())