Notebook Testing Reference

The example notebooks in the documentation are a key touchpoint for getting started and figuring out how to use Toponymy. That means that we want to treat drift between the examples and the code as bugs and catch the drift through testing where possible.

There are a couple of general approaches to testing notebooks:

1. use a notebook runner to run the notebook (e.g. nbclient, nbval)
2. copy all the code out of a notebook and run it as a script (e.g. pytest-nb-as-test)

The advantage of 1) is testing your actual code directly, and keeping dependencies low and standard as everything you need is already required if you're going to be running notebooks anyway. The advantage of 2) is that it sidesteps all the awkwardness of running notebooks in a kernel separate from the test environment so that usual test mocking patterns don't apply and logs/messages don't pass through. In the Toponymy case, choosing 1) based on a tried and true library with lightweight workarounds seemed like the better way to go.

This test suite uses a decorator @notebook_test_replacement(replacement) to fallback on the replacement function when the environment variable NOTEBOOK_TESTING is "true". This means that we can mock and keep the notebook clean of test fallbacks.

We'll see that in the end a test will look like this:

@pytest.mark.parametrize(
	"notebook", active_notebook_list(TEST_NOTEBOOKS, has_openainamer=False)
)

def test_doc_notebook_no_openainamer(notebook, notebook_testing_env):
	cfg = get_notebook_cfg(notebook)
	run_notebook(
		notebook,
		timeout=cfg["timeout"],
	)

There are 4 different layers to getting notebook testing to work using this approach

• A notebook runner
• The notebook test fallback infrastructure
• The notebook discovery
• The CI setup

Let's break each of them down.

Notebook Runner

First, there's run_notebook in toponymy/tools/notebook_runner.py, which forms the bones of all of the test infrastructure.

def run_notebook(
	path: str,
	timeout: int = 3000,
	kernel_name: str = "toponymy-uv",
	instrumented: bool = False,
	return_log_lines: bool = False,
	ignore_litellm: bool = True,
	) -> nbformat.NotebookNode | list[tuple[str, str]]:
	"""
	Execute a Jupyter notebook with optional logging instrumentation and log-line collection.
	

	Runs the notebook in a kernel, injects a logging capture cell to ensure stdlib logging
	is routed to stdout, and optionally collects log-like lines from cell outputs. On exception,
	partial notebook outputs and collected logs are re-emitted before re-raising.
	"""

Because notebooks run in a different kernel than where they are invoked, run_notebook captures any log-like output and passes it on.

For convenience, we also make run_notebook accessible as a script via the command line. In particular,

uv run toponymy/tools/notebook_runner.py doc/basic_usage.ipynb --instrument

• Uses a toponymy-uv kernel by default — register it with uv run python -m ipykernel install --user --name toponymy-uv, or pass a different kernel with --kernel-name.
• --instrument prints per-cell progress and timing.
• Requires the environment to be synced with uv sync --extra example-notebooks.

Dev bonus: running a notebook manually with NOTEBOOK_TESTING=true gets you the same mocked behaviour as the test suite which can be useful for quickly checking if a notebook still works without hitting real APIs or full datasets:

NOTEBOOK_TESTING=true uv run toponymy/tools/notebook_runner.py doc/basic_usage.ipynb

Warning: Don't globally set the NOTEBOOK_TESTING variable to 'true' unless you want fallback behaviour all of the time which is not recommended.

Test Fallback Infrastructure

The test fallback infrastructure are the bones that make @notebook_test_replacement a useful mocking alternative via pytest, and all of the helpers that are invoked by @notebook_test_replacement to mock LLM Namers, shrink data, and handle I/O.

Recall that the fallback decorator @notebook_test_replacement(replacement) routes through the replacement function instead of the decorated function when NOTEBOOK_TESTING is true.

import os
import functools

def notebook_test_replacement(replacement):
	"""
	Decorator for replacing function implementations during example notebook testing.

	When the NOTEBOOK_TESTING environment variable is set to "true", the decorated function
	is replaced with the provided replacement function. This allows lightweight mocking
	during automated notebook runs.
	
	Parameters
	----------
	replacement : callable
		The replacement function to call when NOTEBOOK_TESTING is enabled.

	Returns
	-------
	callable
		A decorator that wraps the target function and applies the replacement conditionally.
	"""
	def decorator(func):
		@functools.wraps(func)
		def wrapper(*args, **kwargs):
			if os.getenv("NOTEBOOK_TESTING", "").lower() == "true":
				return replacement(*args, **kwargs)
			return func(*args, **kwargs)
		return wrapper
	return decorator

To do this cleanly in tests, the pytest fixture notebook_testing_env sets and unsets NOTEBOOK_TESTING, along with other notebook testing related environment variables on a per-test basis.

@pytest.fixture(scope="function")
def notebook_testing_env(notebook_output_dir):
	old = os.environ.get("NOTEBOOK_TESTING")
	old_openai = os.environ.get("OPENAI_API_KEY")
	old_output_dir = os.environ.get("NB_TEST_OUTPUT_DIR")
	
	os.environ["NOTEBOOK_TESTING"] = "true"
	os.environ["OPENAI_API_KEY"] = "notarealkey"
	os.environ["NB_TEST_OUTPUT_DIR"] = str(notebook_output_dir)
	
	try:
		yield
	finally:
		if old is None:
			os.environ.pop("NOTEBOOK_TESTING", None)
		else:
			os.environ["NOTEBOOK_TESTING"] = old
	
		if old_openai is None:
			os.environ.pop("OPENAI_API_KEY", None)	
		else:
			os.environ["OPENAI_API_KEY"] = old_openai	
		if old_output_dir is None:
			os.environ.pop("NB_TEST_OUTPUT_DIR", None)	
		else:
			os.environ["NB_TEST_OUTPUT_DIR"] = old_output_dir

This decorator approach was inspired by https://mstruwig.com/posts/quick-mocking-jupyter-notebook/, but differs in one important way. The reference version (if_testing_mock_with) checks the environment variable at decoration time, which happens when the module is imported and permanently binds the name to either the real function or the replacement. That doesn't work with our notebook_testing_env fixture, which sets/unsets NOTEBOOK_TESTING per test. By the time any test runs, the module has already been imported and the check has already happened. This version moves the check inside the inner wrapper, so it's re-evaluated on every call. This allows notebook_testing_env to control test behaviour.

All this is so that we can use fallbacks instead of the usual pytest mocks. Let's look at the different categories of fallbacks.

Mocking OpenAINamer

For example purposes, we've chosen to use OpenAINamer as the exemplar for LLM namers in the notebooks (although many more are available). However, calling OpenAI during notebook tests isn't ideal and we'd rather fall back on a local model, especially in CI since the tests are meant to test whether a notebook runs to completion and not the connection with an LLM provider.

@notebook_test_replacement(NotebookOpenAINamerMock)
def OpenAINamer(
	model: str = "openai/gpt-4o-mini",
	api_key: str | None = None,
	api_base: str | None = None,
	llm_specific_instructions: str | None = None,	
	max_tokens_topic_name: int = 128,	
	max_tokens_cluster_names: int = 1024,	
	provider_kwargs: dict[str, Any] | None = None,	
	callback: DebugCallback | None = None,	
	base_url: str | None = None, # deprecated, renamed to api_base	
	http_client: "httpx.Client | None" = None, # deprecated, pass via provider_kwargs instead

) -> LiteLLMNamer:

In particular, when the fallback is invoked, OpenAINamer calls are replaced with NotebookOpenAINamerMock calls, which is essentially OllamaNamer under the hood.

def NotebookOpenAINamerMock(*args, **kwargs):
	"""
	For mocking OpenAINamer calls with a local Ollama model.
	"""
	logger.info("Using NotebookOpenAINamerMock instead of OpenAINamer")
	kwargs.pop("base_url", None)
	kwargs.pop("http_client", None)
	kwargs.pop("model", None)
	kwargs.pop("temperature_override", None)
	return OllamaNamer(
		model=get_test_ollama_model(), temperature_override=0.0, **kwargs
	)

Since OllamaNamer is already configured and used for end-to-end integration tests, it's the provider that's also used for notebook testing. In toponymy/tools/notebook_test_helpers.py, the Ollama models for testing are set to match what will be looked for in CI and locally.

Dataset Shrinking

While we want the example notebooks to have big enough datasets to provide a decent example, testing that a notebook runs to completion doesn't need the full dataset, and run time can be dramatically reduced by subsampling the datasets to a smaller size for tests.

We can use the @notebook_test_replacement functionality by writing data loading functions for the datasets that are used in example notebooks. The data loading helpers are found in toponymy/tools/notebook_data_load.py. For instance, in basic_usage.ipynb, we load the 20 newsgroups via:

from toponymy.tools.notebook_data_load import load_newsgroups
newsgroups_df = load_newsgroups()

where

@notebook_test_replacement(load_small_newsgroups)
def load_newsgroups(use_small: bool = False) -> pd.DataFrame:

so that load_newsgroups() will load the 20 newsgroups dataset found at "hf://datasets/lmcinnes/20newsgroups_embedded/data/train-00000-of-00001.parquet", but when the decorator fallback is invoked, will run load_small_newsgroups instead which replaces the datasets with a dataframe of size 150.

Some of the tests don't require data shrinking (they are fast enough and may need the whole dataset to pass), for instance, in clustering_options.ipynb the test already runs in less than a minute so in that case the dataset is loaded via

from toponymy.tools.notebook_data_load import load_newsgroups
newsgroups_df = load_newsgroups(use_small=False)

which loads the full dataset (even though technically it falls back).

There is currently loading functionality for:

• 20 newsgroups: load_newsgroups
• Arxiv Machine Learning: load_arxiv_ml
• Arxiv Category Theory: load_arxiv_ct
• Arxiv Comp Sci: load_bundled_arxiv (the data is bundled into the library as part of the examples directory)

Other datasets can be added similarly as needed.

Redirecting I/O

Finally, a couple of notebooks save/load and as such have side effects that we'd like to control and clean up after tests. For this we create a notebook_output_dir in toponymy/tools/notebook_data_load.py.

Then when doing I/O, as in saving_loading.ipynb, we can use the path to control the location of the file:

topic_model.to_file(notebook_output_dir() /'20ng-topicmodel.tm.zip')

The notebook_output_dir() defaults to the location of the notebook (so the default behaviour if the path was not there), and falls back to _test_output_dir otherwise.

In order to use pytest to setup and teardown the temp directory location, the notebook_testing_env fixture sets NB_TEST_OUTPUT_DIR to a session-wide temporary directory for this purpose. When NB_TEST_OUTPUT_DIR is set, _test_output_dir() will return that value.

Here's what notebook_output_dir() will return:

• By default the notebook's directory.
• With NOTEBOOK_TESTING set to true:
  • If NB_TEST_OUTPUT_DIR is set, then is returns that value. (This is the path when notebook_testing_env is invoked).
  • Else, it creates and tears down a temporary directory for this purpose. (This may happen when using the notebook runner with NOTEBOOK_TESTING set to true).

Notebook Discovery and Config

In toponymy/tools/notebook_test_helpers.py, we have a get_notebooks helper which automatically discovers all of the notebook files in a given directory without xxx in the name. In test_doc_notebooks.py, this is used to generate the list of TEST_NOTEBOOKS from the doc directory, to be tested.

Each notebook has a config entry in NOTEBOOK_CONFIG (in test_doc_notebooks.py) which dictates the following:

• has_openainamer (boolean): whether the notebook uses OpenAINamer (and requires that Ollama be setup and installed for the fall back)
• run_in_pr (boolean): whether to run the notebook on all PRs
• timeout(seconds): a notebook specific timeout for the run based on its usual test run time in CI If a notebook is discovered that does not have a config entry, it gets the very conservative {"has_openainamer": True, "run_in_pr": True, "timeout": 6000}.

Finally, it's all put together into the tests which are parametrized by the doc notebooks:

@pytest.mark.parametrize(
	"notebook", active_notebook_list(TEST_NOTEBOOKS, has_openainamer=False)
)
def test_doc_notebook_no_openainamer(notebook, notebook_testing_env):
	cfg = get_notebook_cfg(notebook)
	
	run_notebook(
		notebook,
		timeout=cfg["timeout"],
	)


@pytest.mark.parametrize(
	"notebook", active_notebook_list(TEST_NOTEBOOKS, has_openainamer=True)
)
def test_doc_notebook_has_openainamer(notebook, notebook_testing_env, ollama_running):
	cfg = get_notebook_cfg(notebook)

	model = get_test_ollama_model()
	logger.info(f"ollama running:{ollama_running}")
	logger.info(f"ollama_has_model {model}:{ollama_has_model(model)}")
	if not ollama_has_model(model):
		pytest.skip(f"{model} not available in local Ollama for OpenAI mocking")
	run_notebook(
		notebook,
		timeout=cfg["timeout"],
	)

Here, active_notebook_list filters down the discovered notebooks TEST_NOTEBOOKS down to the ones required for the test based on whether the trigger was a PR or not and whether the notebook's config has_openainamer. Note that get_notebook_cfg simply returns a notebook's config, and defaults if a notebook has no configs, so that timeouts can be passed through to the runner.

CI Layer

Lastly, we have the CI layer, where notebooks are tested on different triggers. The tests are set up in CI to run on the following triggers:

• On PRs: Always run the notebooks with run_in_pr set to True. Also run any notebooks that have been changed in the PR.
• On a weekly schedule: Scheduled run of all of the doc notebooks once a week.

To detect the changed notebooks on a given PR, there is a DetectChangedNotebooks job in azure-pipeline.yml which sets the environment variable CHANGED_NOTEBOOKS to the list of changed notebooks which gets picked up in test_doc_notebooks.py's should_run_in_pr to add those notebooks to the test suite.

The original PR for adding this infrastructure was in PR #162.