feat: notebook tester + CLI + report summary #924
Conversation
github-actions
bot
added
the
status:awaiting review
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @andycandy, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces a new, automated testing framework for Jupyter notebooks. Its primary goal is to streamline the process of verifying notebook functionality and detecting regressions in their outputs, especially when those outputs might change subtly. The system provides a flexible command-line interface for execution and generates detailed, AI-augmented reports to help developers quickly understand the impact of changes.
Highlights
- Automated Notebook Testing Framework: Introduces a comprehensive system for automatically discovering, executing, and validating Jupyter notebooks.
- Command-Line Interface (CLI): A new cookbook shell script provides a user-friendly entrypoint to run tests, with options for specific notebooks, AI comparison, timeouts, and kernels.
- AI-Powered Output Comparison: Integrates Gemini AI to intelligently compare notebook cell outputs, classifying changes as "ok_cells", "slightly_changed", or "wrong" for regression detection.
- Detailed JSON Reporting: Generates structured JSON reports for each notebook test run, summarizing status, duration, and AI comparison results.
- Colab Compatibility: Includes logic to patch google.colab.userdata.get calls, enabling seamless local execution of notebooks originally designed for Google Colab.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a powerful automated testing framework for Jupyter notebooks, including a shell entrypoint and a Python test runner. The features like automatic discovery, Colab patching, and optional AI-based output comparison are great additions. My review focuses on improving robustness, cross-platform compatibility, and adherence to the project's style guide. I've identified a few issues, including a bug in handling multiple notebook arguments in the cookbook script, some cross-platform compatibility problems, and a few areas in the Python script that could be more robust.
| for a in "$@"; do | ||
| case "$a" in | ||
| --ai-compare) AI_COMPARE=1 ;; | ||
| --timeout=*) TIMEOUT="${a#*=}" ;; | ||
| --kernel=*) KERNEL="${a#*=}" ;; | ||
| *.ipynb) NB_ARG="$a" ;; | ||
| esac | ||
| done |
There was a problem hiding this comment.
The current argument parsing logic for notebook files is brittle. It only handles the last .ipynb file if multiple are provided as separate arguments (e.g., cookbook test nb1.ipynb nb2.ipynb), and it doesn't correctly handle the documented use case for multiple files in a single string (e.g., cookbook test "nb1.ipynb,nb2.ipynb").
A more robust approach is to treat all non-option arguments as notebook paths, collect them, and then join them into a single string for the NB environment variable. This will correctly handle all intended use cases.
NB_ARGS=()
for a in "$@"; do
case "$a" in
--ai-compare) AI_COMPARE=1 ;;
--timeout=*) TIMEOUT="${a#*=}" ;;
--kernel=*) KERNEL="${a#*=}" ;;
*) NB_ARGS+=("$a") ;;
esac
done
NB_ARG=$(IFS=,; echo "${NB_ARGS[*]}")
|
|
||
| def _ensure_requirements(nb_path): | ||
| try: | ||
| import pipreqsnb |
There was a problem hiding this comment.
The rm command is not cross-platform and will fail on Windows. Please use a Python-native way to delete the file, such as os.remove() or pathlib.Path.unlink().
| import pipreqsnb | |
| req_out.unlink() |
| if ot == "stream": buf.append(out.get("text", "")) | ||
| elif ot in ("execute_result", "display_data"): | ||
| data = out.get("data", {}) |
There was a problem hiding this comment.
The return values in these early-exit conditions are incorrect. The function is expected to return a list of dictionaries, but these lines return tuples, which will cause a TypeError in the calling code. The warning about the missing API key should be printed to the console instead of being returned as data.
| if ot == "stream": buf.append(out.get("text", "")) | |
| elif ot in ("execute_result", "display_data"): | |
| data = out.get("data", {}) | |
| if not diffs: return [] | |
| api_key = os.getenv("GOOGLE_API_KEY") | |
| if not api_key: | |
| print("<AI compare skipped: GOOGLE_API_KEY missing>", file=sys.stderr) | |
| return [] |
| if line.strip().startswith(("from google.colab import userdata", "import google.colab")): continue | ||
| if re.match(r"^\s*import\s+os(\s|$)", line): had_os = True | ||
| lines.append(line) | ||
| src = "\n".join(lines) | ||
| def _sub(m): | ||
| key, default = m.group(2), m.group(3) | ||
| return f"os.getenv('{key}', {default})" if default else f"os.getenv('{key}')" | ||
| if "userdata.get(" in src: | ||
| src2 = _USERDATA_RE.sub(_sub, src) | ||
| if src2 != src and not had_os: src = "import os\n" + src2 | ||
| else: src = src2 |
There was a problem hiding this comment.
The img_count variable is initialized but never incremented. The function doesn't seem to be counting images in the cell outputs. To count images, you should check for mime types like image/png, image/jpeg, etc., within the data dictionary of execute_result or display_data outputs.
| if line.strip().startswith(("from google.colab import userdata", "import google.colab")): continue | |
| if re.match(r"^\s*import\s+os(\s|$)", line): had_os = True | |
| lines.append(line) | |
| src = "\n".join(lines) | |
| def _sub(m): | |
| key, default = m.group(2), m.group(3) | |
| return f"os.getenv('{key}', {default})" if default else f"os.getenv('{key}')" | |
| if "userdata.get(" in src: | |
| src2 = _USERDATA_RE.sub(_sub, src) | |
| if src2 != src and not had_os: src = "import os\n" + src2 | |
| else: src = src2 | |
| def _summarize_outputs(outputs): | |
| buf, img_count, err = [], 0, None | |
| for out in outputs or []: | |
| ot = out.get("output_type") | |
| if ot == "stream": buf.append(out.get("text", "")) | |
| elif ot in ("execute_result", "display_data"): | |
| data = out.get("data", {}) | |
| if any(k.startswith("image/") for k in data): | |
| img_count += 1 | |
| text = data.get("text/plain") or "" | |
| buf.append("".join(text) if isinstance(text, list) else str(text)) | |
| elif ot == "error": err = {"ename": out.get("ename"), "evalue": out.get("evalue")} | |
| return {"text": "".join(buf).strip(), "images": img_count, "error": err} |
| """) | ||
|
|
||
|
|
||
| def _coerce_json(text): | ||
| t = text.strip() | ||
| if t.startswith("```"): | ||
| t = t.strip("`") | ||
| if "\n" in t: t = t.split("\n", 1)[1] | ||
| start, end = t.find("["), t.rfind("]") | ||
| if start != -1 and end != -1 and end > start: | ||
| t = t[start:end+1] | ||
| t = re.sub(r",\s*(\]|\})", r"\1", t) | ||
| return json.loads(t) | ||
|
|
||
| results, raw_texts = [], [] | ||
| total_batches = (len(diffs) + batch_size - 1) // batch_size | ||
| for i in range(0, len(diffs), batch_size): | ||
| if progress_callback: | ||
| progress_callback(i // batch_size + 1, total_batches) | ||
|
|
||
| batch = diffs[i:i+batch_size] | ||
| blocks = [f"Cell {d['index']}\n```python\n{d['code'] or ''}\n```\n" | ||
| f"OLD OUTPUT:\n{d['old_text'] or '(empty)'}\n\n" | ||
| f"NEW OUTPUT:\n{d['new_text'] or '(empty)'}\n----\n" for d in batch] | ||
|
|
||
| payload = { | ||
| "system_instruction": {"parts": [{"text": system_text}]}, | ||
| "contents": [{"role": "user", "parts": [{"text": f"File: {file_name}\nEvaluate these cells:\n\n{''.join(blocks)}"}]}], | ||
| "generationConfig": {"temperature": 0.1, "maxOutputTokens": 8192, "response_mime_type": "application/json"} | ||
| } | ||
|
|
||
| try: |
There was a problem hiding this comment.
The raw_texts variable is populated when an exception occurs during the AI comparison, but it's a local variable that is never returned or used. This means exceptions from the API call are silently swallowed. These errors should be logged or propagated so they are visible.
results = []
total_batches = (len(diffs) + batch_size - 1) // batch_size
for i in range(0, len(diffs), batch_size):
if progress_callback:
progress_callback(i // batch_size + 1, total_batches)
batch = diffs[i:i+batch_size]
blocks = [f"Cell {d['index']}\n```python\n{d['code'] or ''}\n```\n"
f"OLD OUTPUT:\n{d['old_text'] or '(empty)'}\n\n"
f"NEW OUTPUT:\n{d['new_text'] or '(empty)'}\n----\n" for d in batch]
payload = {
"system_instruction": {"parts": [{"text": system_text}]},
"contents": [{"role": "user", "parts": [{"text": f"File: {file_name}\nEvaluate these cells:\n\n{''.join(blocks)}"}]}],
"generationConfig": {"temperature": 0.1, "maxOutputTokens": 8192, "response_mime_type": "application/json"}
}
try:
resp = requests.post(url, json=payload, timeout=90)
resp.raise_for_status()
data = resp.json()
part = data.get("candidates", [{}])[0].get("content", {}).get("parts", [{}])[0]
parsed = part if 'text' not in part else _coerce_json(part.get('text', '[]'))
for item in parsed:
idx = int(item.get("index"))
raw_bucket = (item.get("bucket") or "").strip().lower()
bucket = ("ok_cells" if raw_bucket in ("ok", "same", "almost_same", "ok_cells") else
"slightly_changed" if raw_bucket in ("slightly", "slightly_changed") else "wrong")
results.append({"index": idx, "bucket": bucket, "note": (item.get("note") or "").strip()})
except Exception as e:
print(f"<AI compare error: {e}>", file=sys.stderr)| sys.stdout.write(status_line) | ||
|
|
||
| if self.current_file: | ||
| sys.stdout.write(f"Current: {self.current_file}\nStep: {self.current_step}\n\n") |
There was a problem hiding this comment.
This line exceeds the 100-character limit specified in the style guide.1 Please reformat it for better readability.
ok = len(buckets.get("ok_cells", {}))
slight = len(buckets.get("slightly_changed", {}))
wrong = len(buckets.get("wrong", {}))Style Guide References
Footnotes
-
Gemini Cookbook Python notebooks Style Guide, Lines 94-97. ↩
|
Work Still Needed
|
Overview
This PR introduces automated testing of Jupyter notebooks within the repository. It adds a shell entrypoint (cookbook) and a Python test runner (test_nbclient.py) that:
Features
google.colab.userdata.getcalls to use environment variables for seamless local runs.ok_cells,slightly_changed, orwrong.reports/*.compare.jsonfor further analysis.Usage
Entrypoint
Use the cookbook script to run notebook tests:
Options
--ai-compare: Enables AI-based output comparison.--timeout=<seconds>: Sets cell execution timeout (default: 900).--kernel=<name>: Specifies Jupyter kernel (default: python3).[notebook.ipynb]: Path to a specific notebook. If omitted, all notebooks are tested.Output
reports/directory, e.g.,reports/examples__Book_illustration.ipynb.compare.json.Example Report
Each report includes:
passed/failed)ok_cells,slightly_changed,wrong)Notes
GOOGLE_API_KEYandGEMINI_API_KEYare set in your environment.