Skip to content

Commit a0468ce

Browse files
danielhanchendanielhanchen
committed
Add staging CI workflows for unslothai#7191 
1 parent 2d9390f commit a0468ce

19 files changed

Lines changed: 2335 additions & 0 deletions

‎.github/scripts/studio_test_kit/README.md‎

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
# studio_test_kit
2+
3+
A small, vendorable cookbook for driving **Unsloth Studio** end-to-end with
4+
Playwright. Drop these files into your project, read them top-to-bottom,
5+
and you'll know how to script chat flows, generate screenshots + videos,
6+
and run pre/post-PR comparisons.
7+
8+
This is **not** a generic library. It is intentionally Studio-flavoured
9+
(model picker selectors, localStorage keys, `/api/auth/login` route,
10+
`unsloth studio -p` CLI, etc.). See [Customisation](#customisation) for
11+
the override points if you're targeting a different SPA.
12+
13+
## What to vendor (15 files)
14+
15+
Copy the whole `studio_test_kit/` directory. No editable install needed
16+
once on `sys.path`.
17+
18+
| # | File | Purpose | Read for |
19+
|---|------|---------|----------|
20+
| 1 | `__init__.py` | Public re-exports | API surface overview |
21+
| 2 | `lifecycle.py` | Install Studio at a given `UNSLOTH_STUDIO_HOME` from any git ref; launch on a port; parse bootstrap password from log; wait on `/healthz`. | Two-port pre/post comparison setup |
22+
| 3 | `auth.py` | Backend JWT login + Playwright init-script builder that seeds the 5 `unsloth_*` localStorage keys. Convenience constructors for Gemini / OpenAI / Anthropic. | Studio's auth + provider model |
23+
| 4 | `ui.py` | Playwright primitives: `open_chat` async context manager (video recording + post-close rename + optional mp4 transcode), `pick_model`, `set_pill`, `send_prompt`, `wait_for_stream`, `wait_for_image`, `wait_for_text`, `extract_data_url`. | The Playwright lessons |
24+
| 5 | `flows.py` | High-level scenarios returning `FlowResult`: `multi_turn_chat`, `image_generation`, `tool_pills`, `vision_upload`. | Composable test flows |
25+
| 6 | `compose.py` | Post-processing: PIL `hstack_images`/`vstack_images`, ffmpeg `hstack_videos`/`webm_to_mp4`. | Pre/post side-by-side composition |
26+
| 7 | `examples/explicit_walkthrough.py` | **Start here.** Fully inline end-to-end flow with every Playwright primitive at the call site and "why" comments. No `flows.*` shortcuts. | Tutorial |
27+
| 8 | `examples/multi_turn.py` | Thin: 4-turn chat via `flows.multi_turn_chat` | Pattern after you know the basics |
28+
| 9 | `examples/image_gen.py` | Thin: Nano Banana image-gen, saves decoded PNG | Pattern |
29+
| 10 | `examples/tools_pills.py` | Thin: Search + Code composer pills | Pattern |
30+
| 11 | `examples/side_by_side.py` | Two `install_studio` + two `launch_studio` + same flow on both + ffmpeg hstack | Pre/post PR comparison |
31+
| 12 | `examples/__init__.py` | Empty marker | — |
32+
| 13 | `_self_test.py` | 13 offline regression tests (no live Studio needed) | CI / verifying the kit didn't bit-rot |
33+
| 14 | `_smoke_ui.py` | Live Playwright integration test against an in-process HTTP server (verifies init-script seeding, screenshot, video flush, mp4 transcode) | CI / smoke after vendoring |
34+
| 15 | `README.md` | This document | — |
35+
36+
## Setup
37+
38+
```bash
39+
pip install playwright httpx pillow
40+
playwright install chromium
41+
42+
# Optional but recommended (used for mp4 transcode and side-by-side video):
43+
apt-get install ffmpeg # or: brew install ffmpeg
44+
```
45+
46+
Verify the vendoring worked:
47+
48+
```bash
49+
python -m studio_test_kit._self_test # 13/13 offline regression
50+
python -m studio_test_kit._smoke_ui # live Playwright smoke
51+
```
52+
53+
## The model: how Studio authentication and providers work
54+
55+
You need to understand this in 60 seconds before the Playwright bits
56+
make sense.
57+
58+
1. **Login is REST, tokens live in localStorage.** Studio exposes
59+
`POST /api/auth/login` taking `{username, password}` and returning
60+
`{access_token, refresh_token}`. The SPA reads those tokens from
61+
localStorage on every request — NOT from cookies. So
62+
`browser_context.add_cookies()` does nothing useful for auth.
63+
`auth.login()` makes the REST call; `auth.seed_init_script()` plants
64+
the tokens in localStorage.
65+
66+
2. **External providers are localStorage-driven, not server-stored.** The
67+
Settings → Connections UI writes three keys:
68+
69+
| Key | Value |
70+
|-----|-------|
71+
| `unsloth_chat_external_providers` | JSON array of `{id, providerType, name, baseUrl, models}` |
72+
| `unsloth_chat_external_provider_keys` | `{providerId: plaintextKey}` |
73+
| `unsloth_chat_connections_enabled` | `"true"` / `"false"` |
74+
75+
Seeding these via `add_init_script` is exactly equivalent to typing
76+
the key into the dialog manually. The frontend then RSA-encrypts the
77+
plaintext key on every request using the public key from `GET
78+
/api/key`, so the plaintext never goes over the wire even though it
79+
sits in localStorage.
80+
81+
3. **`add_init_script` ordering is what makes seeding work.** The init
82+
script runs on every page navigation BEFORE the SPA's own JS. Plant
83+
AFTER `goto` and you race the SPA's bootstrap — sometimes works,
84+
sometimes silently doesn't.
85+
86+
## Quickstart
87+
88+
Assumes Studio is already running on `--port` (use `lifecycle.launch_studio`
89+
to spin it up programmatically — see `examples/side_by_side.py` for that
90+
pattern).
91+
92+
```bash
93+
export GEMINI_API_KEY=AIza...
94+
95+
# Tutorial: every primitive inline + commentary. Read this file first.
96+
python -m studio_test_kit.examples.explicit_walkthrough \
97+
--port 8902 --password 'YourBootstrap!' \
98+
--model gemini-2.5-flash
99+
100+
# 4-turn conversation (uses flows.multi_turn_chat)
101+
python -m studio_test_kit.examples.multi_turn \
102+
--port 8902 --password 'YourBootstrap!' --model gemini-2.5-flash
103+
104+
# Image generation, saves decoded PNG
105+
python -m studio_test_kit.examples.image_gen \
106+
--port 8902 --password 'YourBootstrap!' \
107+
--model gemini-2.5-flash-image \
108+
--prompt 'A red panda eating ramen in the rain'
109+
110+
# Search + Code composer pills
111+
python -m studio_test_kit.examples.tools_pills \
112+
--port 8902 --password 'YourBootstrap!' --model gemini-2.5-flash
113+
114+
# Full pre/post PR comparison: installs Studio twice, drives both, composes
115+
python -m studio_test_kit.examples.side_by_side \
116+
--pre-branch main --post-branch feat/my-thing \
117+
--pre-port 8901 --post-port 8902 \
118+
--model gemini-2.5-flash
119+
```
120+
121+
## Programmatic flow (custom scenario)
122+
123+
If `flows.multi_turn_chat` / `image_generation` / `tool_pills` /
124+
`vision_upload` don't fit your task, compose primitives:
125+
126+
```python
127+
import asyncio
128+
from pathlib import Path
129+
from studio_test_kit.auth import gemini_provider, login, seed_init_script
130+
from studio_test_kit.ui import (
131+
open_chat, pick_model, send_prompt, set_pill, wait_for_stream,
132+
wait_for_image, extract_data_url,
133+
)
134+
135+
136+
async def run():
137+
auth = await login("http://127.0.0.1:8902", "unsloth", "YourBootstrap!")
138+
init = seed_init_script(
139+
auth,
140+
[gemini_provider(api_key="AIza...", models=["gemini-2.5-flash-image"])],
141+
)
142+
async with open_chat(
143+
"http://127.0.0.1:8902",
144+
init_scripts=[init],
145+
video_dir=Path("out/video"),
146+
video_name="my_run",
147+
transcode_mp4=True,
148+
headless=True,
149+
) as sp:
150+
await pick_model(sp, "gemini-2.5-flash-image")
151+
await set_pill(sp, "images", on=True)
152+
await send_prompt(sp, "Draw a red panda")
153+
data_url = await wait_for_image(sp)
154+
raw = await extract_data_url(data_url)
155+
Path("out/red_panda.png").write_bytes(raw)
156+
print(f"video: {sp.video_webm} mp4: {sp.video_mp4}")
157+
158+
159+
asyncio.run(run())
160+
```
161+
162+
`sp.video_webm` / `sp.video_mp4` are populated AFTER the `async with`
163+
exits (Playwright finalises on `context.close()`).
164+
165+
## Pitfalls (read once, save yourself a day)
166+
167+
Every one of these bit me; the kit defends against all of them.
168+
169+
| # | Pitfall | Fix in the kit |
170+
|---|---|---|
171+
| 1 | `wait_until="networkidle"` never fires on Studio because the chat thread holds long-lived SSE/WebSocket; Playwright deadlines. | `open_chat` uses `wait_until="domcontentloaded"` plus an explicit `form:has(textarea) textarea` visibility wait. |
172+
| 2 | `:has-text("gemini-2.5-flash")` substring-matches `gemini-2.5-flash-image`. | `pick_model` uses `get_by_role("option", name=re.compile(rf"^{re.escape(model_id)}$"))`. |
173+
| 3 | A bare `button:has-text("Search")` clicks the LEFT-SIDEBAR Search-history button, not the composer pill. | All UI selectors are scoped: `form:has(textarea) button:has-text(...)`. |
174+
| 4 | Stop button hides BEFORE the final `<img data:image/...>` paints; screenshots taken on stop-hidden miss the image. | `wait_for_image` polls the DOM via `page.evaluate` for an `<img>` whose `src` starts with `data:image/{png,jpeg,webp}` and validates decoded bytes. |
175+
| 5 | `page.type(...)` simulates ~50 ms per keystroke; long prompts take seconds. | `send_prompt` uses `box.fill(...)` (instant paste). |
176+
| 6 | `page.video.path()` read inside the `async with` returns a tempfile that won't survive context close. | Read `sp.video_webm` AFTER the `async with` exits. |
177+
| 7 | Video stalls if you `await page.close()` before `context.close()`. | `open_chat` closes context then browser in a `finally`. |
178+
| 8 | Init script planted AFTER `goto` races the SPA's bootstrap; sometimes works, sometimes silently doesn't. | `seed_init_script` is wired through `context.add_init_script(...)` BEFORE `new_page`/`goto`. |
179+
| 9 | Cookies: Studio reads its JWT from localStorage, not cookies. | `auth.seed_init_script` plants `unsloth_auth_token` / `unsloth_refresh_token` in localStorage. |
180+
| 10 | Playwright's `record_video_dir` writes `page@<hash>.webm`. Parallel tests sharing a video dir can rename each other's recordings. | `open_chat` snapshots pre-existing webms on entry and renames only files that appeared during this session. |
181+
| 11 | Video finalisation skipped if the test body raises (the case you most want a recording). | `open_chat` wraps the playwright lifecycle in `try/finally`. |
182+
| 12 | Shared deadline between password discovery and `/healthz`: a quiet log starves the healthz check and raises a spurious TimeoutError. | `launch_studio` takes independent `password_timeout_s` (default 30s) and `healthz_timeout_s` (default 180s). |
183+
| 13 | Password log regex `[:\s]+` greedily backtracks on `password = secret`, capturing `"="` as the password. | `_PW_RE` uses an explicit `[:=]?` separator with mandatory `\s+` before the value. |
184+
| 14 | `:has-text` substring on `data:image/` matches SVG tracking pixels. | `wait_for_image(mime_prefixes=("data:image/png", "data:image/jpeg", "data:image/webp"))` excludes SVG by default. |
185+
186+
## Things I avoided (and why)
187+
188+
| Avoided | Why |
189+
|---|---|
190+
| WebKit / Firefox launchers | Studio's RSA-encrypt-on-frontend path uses Web Crypto; only reliable on Chromium for now. |
191+
| `wait_for_selector(...)` | Deprecated in Playwright 1.40+. Use `locator(...).wait_for(state=...)`. |
192+
| `expect_response()` for streaming | SSE doesn't end with a discrete `Response`; DOM polling is sturdier. |
193+
| `evaluate_handle` | Returns a JSHandle you must `.dispose()`. `evaluate` returns serialised JSON, which is enough for selector probes. |
194+
| `playwright codegen` selectors | Produces brittle `.locator("...:nth=2")`; hand-written form-scoped role/text predicates are far more durable. |
195+
| Cookie-based auth | Studio reads JWT from localStorage. Cookies do nothing. |
196+
| Page video at `await page.video.path()` mid-block | The path isn't valid until `context.close()`; reading early gives you a tempfile that gets unlinked. |
197+
| `screencast_*` / CDP video | Lower-level, no auto-flush, no rename. Context-level `record_video_dir` is the only path that just works. |
198+
199+
## Customisation
200+
201+
This kit is tuned to **Unsloth Studio**. If you're driving a similar
202+
FastAPI + React chat app with a different shape, here are the override
203+
points:
204+
205+
| What | File:line | How to swap |
206+
|---|---|---|
207+
| Login URL / payload shape | `auth.py:login` | Replace `POST /api/auth/login` with your endpoint; ensure `StudioAuth(access_token, refresh_token, base_url)` is constructible |
208+
| localStorage key names | `auth.py:seed_init_script` (5 keys) | Edit the `payload = {...}` dict; the rest of the kit only consumes `access_token` directly |
209+
| RSA-on-frontend assumption | `auth.py:5-15` docstring | If your app sends keys plain or via header, modify `seed_init_script` to set whatever your SPA reads |
210+
| Composer scope selector | `ui.py:set_pill / send_prompt` use `form:has(textarea) ...` | Change the prefix if your composer isn't an HTML `<form>` |
211+
| Pill names | `ui.py:Pill` literal + `set_pill` map | Add your toggle labels |
212+
| Model picker trigger | `ui.py:pick_model:trigger` | Adjust the data-testid / button text to match your UI |
213+
| Model picker option role | `ui.py:pick_model:option` | If you don't use `role="option"`, change to `role="menuitem"` etc. |
214+
| Chat route path | `ui.py:open_chat` `goto(f"{base_url}/chat")` | Make `chat_path` a parameter and thread through |
215+
| Stop-button selector | `ui.py:wait_for_stream` | Change `aria-label="Stop generating"` to your label |
216+
| `unsloth studio -p` launcher | `lifecycle.py:launch_studio` | Replace the `bin_path studio -p <port>` command with your CLI |
217+
| Studio install layout (`.venv_t5_550` etc.) | `lifecycle.py:_find_unsloth_bin` | Add globs / accept `bin_search_paths` arg |
218+
| Bootstrap password log shape | `lifecycle.py:_PW_RE` | Edit the regex; current accepts `bootstrap/initial/generated password [is] [:=]? value` |
219+
| Healthz path | `lifecycle.py:launch_studio` `/healthz` | Change endpoint string |
220+
221+
A common refactor for a different app: replace just `auth.py` and the
222+
`lifecycle.*` install/launch helpers. `ui.py` / `flows.py` / `compose.py`
223+
should still work if your SPA has a composer textarea and a streaming
224+
stop button.
225+
226+
## Verification (after vendoring)
227+
228+
The two test runners are part of the kit, not external tooling. Run
229+
them to confirm your vendored copy is intact:
230+
231+
```bash
232+
python -m studio_test_kit._self_test
233+
# 13/13 self-tests passed
234+
235+
python -m studio_test_kit._smoke_ui
236+
# OK -- init script primed token; screenshot ~7 KB;
237+
# webm smoke.webm ~7 KB; mp4 smoke.mp4 ~5 KB
238+
```
239+
240+
`_smoke_ui.py` spins a tiny in-process HTTP server, so it works on a
241+
machine without Studio installed; it's how the kit's CI verifies its
242+
own Playwright pipeline.

‎.github/scripts/studio_test_kit/__init__.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
"""Reusable building blocks for driving Unsloth Studio end-to-end.
2+
3+
Modules:
4+
lifecycle -- Install Studio at a chosen UNSLOTH_STUDIO_HOME from any git
5+
ref, launch it on a chosen port, discover the bootstrap
6+
password, wait for /healthz.
7+
auth -- Backend JWT login, plus Playwright init scripts that seed
8+
localStorage with external providers and plaintext API keys
9+
so the SPA picks them up on first page load.
10+
ui -- Playwright Chromium context manager with video recording,
11+
model picker, composer textarea, pill toggles, send/stop
12+
waits, and a real wait_for_image (polls DOM for a
13+
`data:image/png` <img>, not just stop-button absence).
14+
flows -- High-level scenarios: multi_turn_chat, image_generation,
15+
tool_pills (Search / Code), vision_upload.
16+
compose -- PIL hstack/vstack image composition + ffmpeg hstack video
17+
side-by-side, for pre/post-PR comparisons.
18+
19+
See `examples/` for runnable scripts and `README.md` for the full flow.
20+
"""
21+
22+
from .lifecycle import StudioInstall, install_studio, launch_studio # noqa: F401
23+
from .auth import StudioAuth, ProviderSeed, login, seed_init_script # noqa: F401
24+
from .ui import ( # noqa: F401
25+
StudioPage,
26+
open_chat,
27+
pick_model,
28+
set_pill,
29+
send_prompt,
30+
wait_for_stream,
31+
wait_for_image,
32+
)

0 commit comments

Comments
 (0)