From e7becf5f8873533d7981c8ab6082faf4b8b39cb6 Mon Sep 17 00:00:00 2001 From: Miguel Grinberg Date: Fri, 1 Dec 2023 16:43:09 +0000 Subject: [PATCH] nbtest improvements --- notebooks/document-chunking/.nbtest.yml | 8 + notebooks/document-chunking/Makefile | 11 +- ...nbtest.teardown.with-index-pipelines.ipynb | 74 ++++++++ .../with-index-pipelines.ipynb | 13 +- notebooks/search/.nbtest.yml | 11 ++ notebooks/search/Makefile | 23 +-- notebooks/search/_nbtest.setup.ipynb | 93 ++++++++++ .../search/_nbtest.teardown.03-ELSER.ipynb | 56 ++++++ notebooks/search/_nbtest.teardown.ipynb | 73 ++++++++ test/README.md | 41 +---- test/nbtest/README.md | 115 +++++++++++++ test/nbtest/nbtest.py | 160 +++++++++++------- test/nbtest/requirements.in | 1 + test/nbtest/requirements.txt | 4 +- 14 files changed, 559 insertions(+), 124 deletions(-) create mode 100644 notebooks/document-chunking/.nbtest.yml create mode 100644 notebooks/document-chunking/_nbtest.teardown.with-index-pipelines.ipynb create mode 100644 notebooks/search/.nbtest.yml create mode 100644 notebooks/search/_nbtest.setup.ipynb create mode 100644 notebooks/search/_nbtest.teardown.03-ELSER.ipynb create mode 100644 notebooks/search/_nbtest.teardown.ipynb create mode 100644 test/nbtest/README.md diff --git a/notebooks/document-chunking/.nbtest.yml b/notebooks/document-chunking/.nbtest.yml new file mode 100644 index 00000000..7754ba8f --- /dev/null +++ b/notebooks/document-chunking/.nbtest.yml @@ -0,0 +1,8 @@ +masks: +- "'name': '[^']+'" +- "'cluster_name': '[^']+'" +- "'cluster_uuid': '[^']+'" +- "'build_flavor': '[^']+'" +- '[0-9]+\.[0-9]+\.[0-9]+' +- "'build_hash': '[^']+'" +- "'build_date': '[^']+'" diff --git a/notebooks/document-chunking/Makefile b/notebooks/document-chunking/Makefile index ae3f1877..c4831ac5 100644 --- a/notebooks/document-chunking/Makefile +++ b/notebooks/document-chunking/Makefile @@ -1,7 +1,10 @@ NBTEST = ../../bin/nbtest +NOTEBOOKS = \ + with-index-pipelines.ipynb -.PHONY: all +.PHONY: all $(NOTEBOOKS) -all: - $(NBTEST) \ - with-index-pipelines.ipynb +all: $(NOTEBOOKS) + +$(NOTEBOOKS): + $(NBTEST) $@ diff --git a/notebooks/document-chunking/_nbtest.teardown.with-index-pipelines.ipynb b/notebooks/document-chunking/_nbtest.teardown.with-index-pipelines.ipynb new file mode 100644 index 00000000..fd007cf5 --- /dev/null +++ b/notebooks/document-chunking/_nbtest.teardown.with-index-pipelines.ipynb @@ -0,0 +1,74 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "1422b7bb-bc8c-42bb-b070-53fce3cf6144", + "metadata": {}, + "outputs": [], + "source": [ + "from elasticsearch import Elasticsearch\n", + "from getpass import getpass\n", + "\n", + "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n", + "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n", + "\n", + "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n", + "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n", + "\n", + "# Create the client instance\n", + "client = Elasticsearch(\n", + " # For local development\n", + " # hosts=[\"http://localhost:9200\"] \n", + " cloud_id=ELASTIC_CLOUD_ID,\n", + " api_key=ELASTIC_API_KEY,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e4a89367-d23a-4340-bc92-2dcabd18adcd", + "metadata": {}, + "outputs": [], + "source": [ + "client.indices.delete(index=\"chunk_passages_example\")\n", + "client.ingest.delete_pipeline(id=\"chunk_text_to_passages\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4ac37f1b-6122-49fe-a3b8-e8f2025a0961", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " client.ml.delete_trained_model(model_id=\"sentence-transformers__all-minilm-l6-v2\", force=True)\n", + "except:\n", + " pass" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/document-chunking/with-index-pipelines.ipynb b/notebooks/document-chunking/with-index-pipelines.ipynb index 34813ec2..cd5e92fa 100644 --- a/notebooks/document-chunking/with-index-pipelines.ipynb +++ b/notebooks/document-chunking/with-index-pipelines.ipynb @@ -523,17 +523,6 @@ "\n", "pretty_response(response)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "b269da89", - "metadata": {}, - "outputs": [], - "source": [ - "client.indices.delete(index=INDEX_NAME)\n", - "client.ingest.delete_pipeline(id=\"chunk_text_to_passages\")\n" - ] } ], "metadata": { @@ -555,7 +544,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.3" + "version": "3.11.6" } }, "nbformat": 4, diff --git a/notebooks/search/.nbtest.yml b/notebooks/search/.nbtest.yml new file mode 100644 index 00000000..e2c4921f --- /dev/null +++ b/notebooks/search/.nbtest.yml @@ -0,0 +1,11 @@ +masks: +- "'name': '[^']+'" +- "'build_flavor': '[^']+'" +- '[0-9]+\.[0-9]+\.[0-9]+' +- "'cluster_name': '[^']+'" +- "'cluster_uuid': '[^']+'" +- "'build_hash': '[^']+'" +- "'build_date': '[^']+'" +- "'_version': [0-9]+" +- '^ID: .*$' +- '^Score: [0-9]+\.[0-9][0-9]*$' diff --git a/notebooks/search/Makefile b/notebooks/search/Makefile index 5de8acab..21da1353 100644 --- a/notebooks/search/Makefile +++ b/notebooks/search/Makefile @@ -1,13 +1,16 @@ NBTEST = ../../bin/nbtest +NOTEBOOKS = \ + 00-quick-start.ipynb \ + 01-keyword-querying-filtering.ipynb \ + 02-hybrid-search.ipynb \ + 03-ELSER.ipynb \ + 04-multilingual.ipynb \ + 05-query-rules.ipynb \ + 06-synonyms-api.ipynb -.PHONY: all +.PHONY: all $(NOTEBOOKS) -all: - $(NBTEST) \ - 00-quick-start.ipynb \ - 01-keyword-querying-filtering.ipynb \ - 02-hybrid-search.ipynb \ - 03-ELSER.ipynb \ - 04-multilingual.ipynb \ - 05-query-rules.ipynb \ - 06-synonyms-api.ipynb +all: $(NOTEBOOKS) + +$(NOTEBOOKS): + $(NBTEST) $@ diff --git a/notebooks/search/_nbtest.setup.ipynb b/notebooks/search/_nbtest.setup.ipynb new file mode 100644 index 00000000..2d321ae7 --- /dev/null +++ b/notebooks/search/_nbtest.setup.ipynb @@ -0,0 +1,93 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "e180af3a-3a2c-4186-a577-7051ec6460b1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install -qU elasticsearch sentence-transformers" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "63d22ea2-ecca-41bb-b08f-de8ad49cda41", + "metadata": {}, + "outputs": [], + "source": [ + "# get the Elasticsearch client\n", + "from elasticsearch import Elasticsearch\n", + "from getpass import getpass\n", + "\n", + "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n", + "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n", + "\n", + "client = Elasticsearch(cloud_id=ELASTIC_CLOUD_ID, api_key=ELASTIC_API_KEY,)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b367acaa-90e6-43d0-b9ae-cf42a0e2c0f1", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from urllib.request import urlopen\n", + "from sentence_transformers import SentenceTransformer\n", + "\n", + "if NBTEST[\"notebook\"] in ['01-keyword-querying-filtering.ipynb', '02-hybrid-search.ipynb', '06-synonyms-api.ipynb']:\n", + " # these tests need book_index to exist ahead of time\n", + " client.indices.delete(index=\"book_index\", ignore_unavailable=True)\n", + " \n", + " mappings = {\n", + " \"properties\": {\n", + " \"title_vector\": {\n", + " \"type\": \"dense_vector\",\n", + " \"dims\": 384,\n", + " \"index\": \"true\",\n", + " \"similarity\": \"cosine\"\n", + " }\n", + " }\n", + " }\n", + " client.indices.create(index='book_index', mappings=mappings)\n", + "\n", + " url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/notebooks/search/data.json\"\n", + " response = urlopen(url)\n", + " books = json.loads(response.read())\n", + "\n", + " model = SentenceTransformer('all-MiniLM-L6-v2')\n", + " operations = []\n", + " for book in books:\n", + " operations.append({\"index\": {\"_index\": \"book_index\"}})\n", + " # Transforming the title into an embedding using the model\n", + " book[\"title_vector\"] = model.encode(book[\"title\"]).tolist()\n", + " operations.append(book)\n", + " client.bulk(index=\"book_index\", operations=operations, refresh=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/search/_nbtest.teardown.03-ELSER.ipynb b/notebooks/search/_nbtest.teardown.03-ELSER.ipynb new file mode 100644 index 00000000..43276f1f --- /dev/null +++ b/notebooks/search/_nbtest.teardown.03-ELSER.ipynb @@ -0,0 +1,56 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "7bf006aa-91cf-4c3a-b685-1f8ca5892a33", + "metadata": {}, + "outputs": [], + "source": [ + "from elasticsearch import Elasticsearch\n", + "from getpass import getpass\n", + "\n", + "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n", + "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n", + "\n", + "client = Elasticsearch(cloud_id=ELASTIC_CLOUD_ID, api_key=ELASTIC_API_KEY,)\n", + "\n", + "# delete the notebook's index\n", + "client.indices.delete(index=\"elser-example-movies\", ignore_unavailable=True)\n", + "\n", + "# delete the pipeline\n", + "try:\n", + " client.ingest.delete_pipeline(id=\"elser-ingest-pipeline\")\n", + "except:\n", + " pass\n", + "\n", + "# delete the model\n", + "try:\n", + " client.ml.delete_trained_model(model_id=\".elser_model_2\", force=True)\n", + "except:\n", + " pass" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/search/_nbtest.teardown.ipynb b/notebooks/search/_nbtest.teardown.ipynb new file mode 100644 index 00000000..c66a543a --- /dev/null +++ b/notebooks/search/_nbtest.teardown.ipynb @@ -0,0 +1,73 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "7bcf0f81-aec8-4f49-918c-3163917885ec", + "metadata": {}, + "outputs": [], + "source": [ + "indexes = {\n", + " \"00-quick-start.ipynb\": \"book_index\",\n", + " \"01-keyword-querying-filtering.ipynb\": \"book_index\",\n", + " \"02-hybrid-search.ipynb\": \"book_index\",\n", + " # 03-ELSER.ipynb has its own teardown notebook\n", + " \"04-multilingual.ipynb\": \"articles\",\n", + " \"05-query-rules.ipynb\": \"products_index\",\n", + " \"06-synonyms-api.ipynb\": \"book_index\",\n", + "}\n", + "INDEX_NAME = indexes.get(NBTEST[\"notebook\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fcd17ce3-ece3-4268-b37b-bbf47c2437c8", + "metadata": {}, + "outputs": [], + "source": [ + "# get the Elasticsearch client\n", + "from elasticsearch import Elasticsearch\n", + "from getpass import getpass\n", + "\n", + "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n", + "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n", + "\n", + "client = Elasticsearch(cloud_id=ELASTIC_CLOUD_ID, api_key=ELASTIC_API_KEY,)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "abf51067-61f8-4cf3-b950-464805ea0e8d", + "metadata": {}, + "outputs": [], + "source": [ + "# delete the notebook's index\n", + "if INDEX_NAME:\n", + " client.indices.delete(index=INDEX_NAME, ignore_unavailable=True)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/test/README.md b/test/README.md index 5b6ae30e..d08da88e 100644 --- a/test/README.md +++ b/test/README.md @@ -1,41 +1,4 @@ # Testing -The `nbtest` directory contains a tool that can validate Python notebooks. A -wrapper script to install and run this tool is in the `bin/nbtest` directory in -this repository. - -Example usage: - -```bash -bin/nbtest notebook/search/00-quick-start.ipynb -``` - -To test all notebooks, you can run `make` from the top-level directory: - -```bash -make -``` - -## Handling of `getpass` - -The `nbtest` script runs the notebooks with an alternative version of the -`getpass()` function that looks for requested values in environment variables -that need to be set before invoking the script. - -Consider the following example, which is used in many Elastic notebooks: - -```python -CLOUD_ID = getpass("Elastic Cloud ID:") -ELASTIC_API_KEY = getpass("Elastic Api Key:") -``` - -The `getpass()` function used by `nbtest` takes the prompt given as an -argument, and converts it to an environment variable name with the following -rules: - -- Spaces are converted to underscores -- Non-alphanumeric characters are removed -- Letters are uppercased - -In the above example, the variables that will be used are `ELASTIC_CLOUD_ID` -and `ELASTIC_API_KEY`. +This directory contains tooling to help us test and maintain the Search Labs +Python notebooks. diff --git a/test/nbtest/README.md b/test/nbtest/README.md new file mode 100644 index 00000000..52629356 --- /dev/null +++ b/test/nbtest/README.md @@ -0,0 +1,115 @@ +# `nbtest` + +The `nbtest` directory contains a tool that can validate Python notebooks. A +wrapper script to install and run this tool is in the `bin/nbtest` directory in +this repository. + +Example usage: + +```bash +bin/nbtest notebook/search/00-quick-start.ipynb +``` + +To test all notebooks that are supported under this tool, you can run `make` +from the top-level directory: + +```bash +make +``` + +## How it works + +`nbtest` runs all the code cells in the notebook in order from top to bottom, +and reports two error situations: + +- If any code cells raise an unexpected exception +- If any code cells that have output saved in the notebook generate a different output (not counting specially designated sections, see the "Configuration" section below) + +Something to keep in mind when designing notebooks that are testable is that +for any operations that are asynchronous it is necessary to add code that +blocks until these operations complete, so that the entire notebook can +execute in batch mode without errors. + +## Configuration + +`nbtest` looks for a configuration file named `.nbtest.yml` in the same +directory as the target notebook. If the configuration file is found, it is +imported and applied while the notebook runs. + +There is currently one supported configuration variable, called `masks`. This +variable can be set to a list of regular expresssions for details in the output +of cells that are variable in nature and should be masked when comparing the +previously stored output against output from the current run. + +Here is an example `.nbtest.yml` file: + +```yaml +masks: +- "'name': '[^']+'" +- "'cluster_name': '[^']+'" +- "'cluster_uuid': '[^']+'" +- "'build_flavor': '[^']+'" +- '[0-9]+\.[0-9]+\.[0-9]+' +- "'build_hash': '[^']+'" +- "'build_date': '[^']+'" +``` + +## Handling of `getpass` + +The `nbtest` script runs the notebooks with an alternative version of the +`getpass()` function that looks for requested values in environment variables +that need to be set before invoking the script. + +Consider the following example, which is used in many Elastic notebooks: + +```python +CLOUD_ID = getpass("Elastic Cloud ID:") +ELASTIC_API_KEY = getpass("Elastic Api Key:") +``` + +The `getpass()` function used by `nbtest` takes the prompt given as an +argument, and converts it to an environment variable name with the following +rules: + +- Spaces are converted to underscores +- Non-alphanumeric characters are removed +- Letters are uppercased + +In the above example, the variables that will be used are `ELASTIC_CLOUD_ID` +and `ELASTIC_API_KEY`. + +## Set up and tear down procedures + +Sometimes it is necessary to perform "set up" and/or "tear down" operations +before and after a notebook runs. `nbtest` will look for notebooks with special +names designated as set up or tear down and execute those notebooks to perform +any necessary code. + +Note that if errors occur while executing a set up or tear down notebook, the +errors are reported, but are not counted as a test failure. + +### Set up notebooks + +`nbtest` will look for the following notebooks names and execute any that are +found before running the target notebook: + +- `_nbtest_setup.ipynb` +- `_nbtest_setup.[notebook-name].ipynb` + +The generic notebook can be used for general set up logic that applies to more +than one notebook in the same directory. Inside these notebooks, the +`NBTEST["notebook"]` expression can be used to obtain the name of the notebook +under test. + +### Tear down notebooks + +`nbtest` will look for the following notebooks names and execute any that are +found after running the target notebook, regardless of the testing having +succeeded or failed: + +- `_nbtest_teardown.[notebook-name].ipynb` +- `_nbtest_teardown.ipynb` + +These notebooks are inteded for cleanup that needs to happen after a text, for +example to delete indexes. As in the set up case, `NBTEST["notebook"]` is set +to the notebook that was tested. diff --git a/test/nbtest/nbtest.py b/test/nbtest/nbtest.py index 40b91125..0dbe492c 100755 --- a/test/nbtest/nbtest.py +++ b/test/nbtest/nbtest.py @@ -5,6 +5,7 @@ import os import re import sys +import yaml # these suppress jupyter warnings on startup os.environ['PYDEVD_DISABLE_FILE_VALIDATION'] = '1' @@ -13,26 +14,23 @@ from jupyter_core.paths import jupyter_data_dir from jupyter_client.kernelspec import KernelSpecManager import nbformat -from nbconvert.preprocessors import ExecutePreprocessor +from nbconvert.preprocessors import ExecutePreprocessor as OriginalExecutePreprocessor from rich import print as rprint from rich.markdown import Markdown basedir = os.path.abspath(os.path.dirname(__file__)) -# The list of tuples below is used for masking parts of the output of code -# cells so that they do not trigger spurious diff errors. -# Each tuple has a regular expression to match, and a replacement string. -MASKS = [ - (r'instance-[0-9]+', 'instance-XXXXX'), - (r'[0-9]+\.[0-9]+\.[0-9]+', '..'), - (r"'cluster_name': '[^']+'", "'cluster_name': 'XXXXX'"), - (r"'cluster_uuid': '[^']+'", "'cluster_uuid': 'XXXXX'"), - (r"'build_hash': '[^']+'", "'build_hash': 'XXXXX'"), - (r"'build_date': '[^']+'", "'build_date': 'XXXXX'"), - (r"'_version': [0-9]+", "'_version': X"), - (r'^ID: .*$', 'ID: XXXXX', 'ID: XXXXXX'), - (r'^Score: [0-9]+\.[0-9][0-9]*$', 'Score: X.XXX'), -] + +class ExecutePreprocessor(OriginalExecutePreprocessor): + def __init__(self, **kwargs): + self.inject = kwargs.pop('inject', {}) + super().__init__(**kwargs) + + def preprocess_cell(self, cell, resources, index): + """Inject the notebook name as variable NBTEST_NOTEBOOK""" + if cell['cell_type'] == 'code': + cell['source'] = f'NBTEST = {self.inject}\n{cell["source"]}' + return super().preprocess_cell(cell, resources, index) def register_python3_test_kernel(): @@ -46,10 +44,10 @@ def unregister_python3_test_kernel(): kernel_spec_manager.remove_kernel_spec('python3-test') -def preprocess_output(output): +def preprocess_output(output, masks): """This function masks the output to hide insignificant differences.""" - for mask in MASKS: - output = re.sub(mask[0], mask[1], output, flags=re.MULTILINE) + for mask in masks: + output = re.sub(mask, '', output, flags=re.MULTILINE) return output @@ -62,10 +60,48 @@ def diff_output(source_output, test_output): rprint(Markdown(f'```diff\n{diff}```\n', code_theme='vim')) +def nbtest_setup_teardown(notebooks, inject={}): + ep = ExecutePreprocessor(timeout=600, kernel_name='python3-test', + inject=inject) + for nb in notebooks: + try: + with open(nb, 'rt') as f: + nb = nbformat.read(f, as_version=4) + except FileNotFoundError: + pass + else: + try: + ep.preprocess(nb, {'metadata': {'path': basedir}}) + except Exception as exc: + rprint(f' [red]Failed in {nb}[default]') + print(exc) + return 1 + + def nbtest_one(notebook, verbose): """Run a notebook and ensure output is the same as in the original.""" rprint(f'Running [yellow]{notebook}[default]...', end='') + notebook_dir = os.path.dirname(notebook) + notebook_name = os.path.basename(notebook) + + # import the .nbtest.yml config file from the notebook's directory + config = {'masks': []} + try: + with open(os.path.join(notebook_dir, '.nbtest.yml'), mode='rt') as f: + config.update(yaml.safe_load(f.read())) + except FileNotFoundError: + pass + + # run the setup notebooks (if available) + setup_notebooks = [ + os.path.join(notebook_dir, '_nbtest.setup.ipynb'), + os.path.join(notebook_dir, f'_nbtest.setup.{notebook_name}'), + ] + nbtest_setup_teardown(setup_notebooks, inject={'notebook': notebook_name}) + + # run the target notebook + ep = ExecutePreprocessor(timeout=600, kernel_name='python3-test') try: with open(notebook, 'rt') as f: nb = nbformat.read(f, as_version=4) @@ -74,52 +110,59 @@ def nbtest_one(notebook, verbose): return 1 original_cells = deepcopy(nb.cells) - ep = ExecutePreprocessor(timeout=600, kernel_name='python3-test') + ret = 0 try: ep.preprocess(nb, {'metadata': {'path': basedir}}) except Exception as exc: rprint(' [red]Failed[default]') print(exc) - return 1 + ret = 1 - ret = 0 - cell = 0 - for source, test in zip(original_cells, nb.cells): - cell += 1 - if source['cell_type'] == 'code': - source_output = { - output.get('name', '?'): output.get('text', '') - for output in source['outputs'] - } - test_output = { - output.get('name', '?'): output.get('text', '') - for output in test['outputs'] - } - for name in source_output: - if name not in ['stdout', 'stderr']: - if verbose: - rprint(f'>>>>> [magenta]code cell #{cell}({name})' - '[default]: [dim white]Skipped[default]') - continue - base = preprocess_output(str(source_output[name])) - current = preprocess_output(str(test_output.get(name, ''))) - if base == current: - if verbose: - rprint(f'>>>>> [yellow]code cell #{cell}/({name})' - '[default]: [green]OK[default]') - else: - if ret == 0: - ret = 1 - rprint(' [red]Failed[default]') - rprint(f'>>>>> [yellow]code cell #{cell}/({name})' - '[default]: [red]Error[default]') - diff_output(base, current) - else: - if verbose: - rprint(f'>>>>> [magenta]{source["cell_type"]} cell #{cell}' - '[default]: [dim white]Skipped[default]') if ret == 0: - rprint(' [green]OK[default]') + cell = 0 + for source, test in zip(original_cells, nb.cells): + cell += 1 + if source['cell_type'] == 'code': + source_output = { + output.get('name', '?'): output.get('text', '') + for output in source['outputs'] + } + test_output = { + output.get('name', '?'): output.get('text', '') + for output in test['outputs'] + } + for name in source_output: + if name not in ['stdout', 'stderr']: + if verbose: + rprint(f'>>>>> [magenta]code cell #{cell}({name})' + '[default]: [dim white]Skipped[default]') + continue + base = preprocess_output(str(source_output[name]), config['masks']) + current = preprocess_output(str(test_output.get(name, '')), config['masks']) + if base == current: + if verbose: + rprint(f'>>>>> [yellow]code cell #{cell}/({name})' + '[default]: [green]OK[default]') + else: + if ret == 0: + ret = 1 + rprint(' [red]Failed[default]') + rprint(f'>>>>> [yellow]code cell #{cell}/({name})' + '[default]: [red]Error[default]') + diff_output(base, current) + else: + if verbose: + rprint(f'>>>>> [magenta]{source["cell_type"]} cell #{cell}' + '[default]: [dim white]Skipped[default]') + if ret == 0: + rprint(' [green]OK[default]') + + # run the teardown notebooks (if available) + teardown_notebooks = [ + os.path.join(notebook_dir, f'_nbtest.teardown.{notebook_name}'), + os.path.join(notebook_dir, '_nbtest.teardown.ipynb'), + ] + nbtest_setup_teardown(teardown_notebooks, inject={'notebook': notebook_name}) return ret @@ -130,7 +173,8 @@ def nbtest(notebook, verbose): """ ret = 0 for nb in notebook: - ret += nbtest_one(notebook=nb, verbose=verbose) + if not nb.startswith('_nbtest'): + ret += nbtest_one(notebook=nb, verbose=verbose) return ret diff --git a/test/nbtest/requirements.in b/test/nbtest/requirements.in index ee3709c6..6182dbde 100644 --- a/test/nbtest/requirements.in +++ b/test/nbtest/requirements.in @@ -1,3 +1,4 @@ pip-tools jupyter rich +pyyaml diff --git a/test/nbtest/requirements.txt b/test/nbtest/requirements.txt index cdd983cf..58cb336e 100644 --- a/test/nbtest/requirements.txt +++ b/test/nbtest/requirements.txt @@ -233,7 +233,9 @@ python-dateutil==2.8.2 python-json-logger==2.0.7 # via jupyter-events pyyaml==6.0.1 - # via jupyter-events + # via + # -r requirements.in + # jupyter-events pyzmq==25.1.1 # via # ipykernel