From 11ec374e88a1cc1532220b008d2f331c65ebbef2 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Tue, 21 May 2024 09:42:56 +0100
Subject: [PATCH] Add per-version testing restrictions

---
 .github/workflows/tests.yml                   |  4 +-
 bin/find-notebooks-to-test.sh                 | 42 ++++++++++++++++++-
 .../with-index-pipelines.ipynb                |  1 +
 .../with-langchain-splitters.ipynb            |  2 +-
 4 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 1ec66a5..c484f0c 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -49,6 +49,6 @@ jobs:
         run: make install-nbtest
       - name: Warm up
         continue-on-error: true
-        run: sleep 30 && PATCH_ES=1 ELASTIC_CLOUD_ID=foo ELASTIC_API_KEY=bar bin/nbtest notebooks/search/00-quick-start.ipynb
+        run: sleep 30 && PATCH_ES=1 ELASTIC_CLOUD_ID=foo ELASTIC_API_KEY=bar ES_STACK=${{ matrix.es_stack }} bin/nbtest notebooks/search/00-quick-start.ipynb
       - name: Run tests
-        run: PATCH_ES=1 FORCE_COLOR=1 make -s test
+        run: PATCH_ES=1 FORCE_COLOR=1 ES_STACK=${{ matrix.es_stack }} make -s test
diff --git a/bin/find-notebooks-to-test.sh b/bin/find-notebooks-to-test.sh
index f2221ff..58c4ce3 100755
--- a/bin/find-notebooks-to-test.sh
+++ b/bin/find-notebooks-to-test.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
-# add any notebooks that are currently not testable to the exempt list
+
 EXEMPT_NOTEBOOKS=(
+    # Add any notebooks that are currently not testable to the exempt list
     "notebooks/esql/esql-getting-started.ipynb"
     "notebooks/search/07-inference.ipynb"
     "notebooks/search/08-learning-to-rank.ipynb"
@@ -25,9 +26,46 @@ EXEMPT_NOTEBOOKS=(
     "notebooks/enterprise-search/app-search-engine-exporter.ipynb"
 )
 
+# Per-version testing exceptions
+# use variables named EXEMPT_NOTEBOOKS__{major}_[minor}_{patch} to list
+# notebooks that cannot run on that stack version or older
+# Examples:
+# EXEMPT_NOTEBOOKS__8 for notebooks that must be skipped on all versions 8.x and older
+# EXEMPT_NOTEBOOKS__8_14 for notebooks that must skipped on versions 8.14 and older
+# EXEMPT_NOTEBOOKS__8_13_2 for notebooks that must be skipped on versions 8.13.2 and older
+
+EXEMPT_NOTEBOOKS__8_13=(
+    # Add any notebooks that must be skipped on versions 8.13 or older here
+    "notebooks/search/03-ELSER.ipynb"
+    "notebooks/model-upgrades/upgrading-index-to-use-elser.ipynb"
+    "notebooks/langchain/langchain-vector-store-using-elser.ipynb"
+)
+
+# this function parses a version given as M[.N[.P]] or M[_N[_P]] into a numeric form
+function parse_version { echo "$@" | awk -F'[._]' '{ printf("%02d%02d%02d\n", $1, $2, $3); }'; }
+
+# this is the version CI is running
+ci_version=$(parse_version ${ES_STACK:-99.99})
+
 ALL_NOTEBOOKS=$(find notebooks -name "*.ipynb" | grep -v "_nbtest" | grep -v ".ipynb_checkpoints" | sort)
 for notebook in $ALL_NOTEBOOKS; do
-    if [[ ! "${EXEMPT_NOTEBOOKS[@]}" =~ $notebook ]]; then
+    skip=
+    # check the master exception list
+    if [[ "${EXEMPT_NOTEBOOKS[@]}" =~ $notebook ]]; then
+        skip=yes
+    else
+        # check the versioned exception lists
+        for exempt_key in ${!EXEMPT_NOTEBOOKS__*}; do
+            exempt_version=$(parse_version ${exempt_key/EXEMPT_NOTEBOOKS__/})
+            if [ $exempt_version -ge $ci_version ]; then
+                exempt_notebooks=$(eval 'echo ${'${exempt_key}'[@]}')
+                if [[ "${exempt_notebooks[@]}" =~ $notebook ]]; then
+                    skip=yes
+                fi
+            fi
+        done
+    fi
+    if [[ "$skip" == "" ]]; then
         echo $notebook
     fi
 done
diff --git a/notebooks/document-chunking/with-index-pipelines.ipynb b/notebooks/document-chunking/with-index-pipelines.ipynb
index 55fbcdc..1ce5bf8 100644
--- a/notebooks/document-chunking/with-index-pipelines.ipynb
+++ b/notebooks/document-chunking/with-index-pipelines.ipynb
@@ -84,6 +84,7 @@
    "source": [
     "from elasticsearch import Elasticsearch\n",
     "from getpass import getpass\n",
+    "import time\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",
diff --git a/notebooks/document-chunking/with-langchain-splitters.ipynb b/notebooks/document-chunking/with-langchain-splitters.ipynb
index 003f612..5b502eb 100644
--- a/notebooks/document-chunking/with-langchain-splitters.ipynb
+++ b/notebooks/document-chunking/with-langchain-splitters.ipynb
@@ -31,7 +31,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!python3 -m pip install -qU langchain langchain-elasticsearch elasticsearch eland==8.12.1 jq"
+    "!python3 -m pip install -qU langchain langchain-community langchain-elasticsearch elasticsearch eland==8.12.1 jq"
    ]
   },
   {