From 7009b411a62f0d9682a69ea9d4103d24f34eaedf Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Tue, 30 Jan 2024 12:40:17 +0000
Subject: [PATCH] add nbtest job to the GitHub Actions CI

---
 .github/workflows/tests.yml                   | 25 +++++++++++++++++++
 bin/mocks.py                                  | 18 +++++++++++++
 bin/nbtest                                    | 10 +++++++-
 .../integrations/hugging-face/.nbtest.yml     |  2 ++
 4 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 .github/workflows/tests.yml
 create mode 100644 bin/mocks.py
 create mode 100644 notebooks/integrations/hugging-face/.nbtest.yml

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 00000000..f8d5971c
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,25 @@
+name: notebook-tests
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    services:
+      elasticsearch:
+        image: docker.elastic.co/elasticsearch/elasticsearch:8.12
+        env:
+          discovery.type: single-node
+          xpack.security.enabled: false
+          xpack.security.http.ssl.enabled: false
+          xpack.license.self_generated.type: trial
+        ports:
+          - 9200:9200
+    steps:
+      - uses: actions/checkout@v4
+      - run: sleep 30
+      - run: MOCK_ES=1 make test
diff --git a/bin/mocks.py b/bin/mocks.py
new file mode 100644
index 00000000..99d9cba4
--- /dev/null
+++ b/bin/mocks.py
@@ -0,0 +1,18 @@
+# This module is provided to nbtest when running under GitHub Actions CI to
+# mock the Elasticsearch client to always connect on localhost, regardless of
+# arguments provided
+
+import elasticsearch
+
+orig_es_init = elasticsearch.Elasticsearch.__init__
+
+
+def patched_es_init(self, *args, **kwargs):
+    if 'cloud_id' in kwargs:
+        assert kwargs['cloud_id'] == 'foo'
+    if 'api_key' in kwargs:
+        assert kwargs['api_key'] == 'bar'
+    return orig_es_init(self, 'http://localhost:9200')
+
+
+elasticsearch.Elasticsearch.__init__ = patched_es_init
diff --git a/bin/nbtest b/bin/nbtest
index 3cd38763..fda0a185 100755
--- a/bin/nbtest
+++ b/bin/nbtest
@@ -4,5 +4,13 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 if [[ ! -f $SCRIPT_DIR/../.venv/bin/nbtest ]]; then
     make nbtest
 fi
+
+NBTEST="nbtest"
+if [[ "$MOCK_ES" != "" ]]; then
+    export ELASTIC_CLOUD_ID=foo
+    export ELASTIC_API_KEY=bar
+    NBTEST="$NBTEST --mocks bin.mocks"
+fi
+
 source $SCRIPT_DIR/../.venv/bin/activate
-$SCRIPT_DIR/../.venv/bin/nbtest $*
+$SCRIPT_DIR/../.venv/bin/$NBTEST $*
diff --git a/notebooks/integrations/hugging-face/.nbtest.yml b/notebooks/integrations/hugging-face/.nbtest.yml
new file mode 100644
index 00000000..0b88e44d
--- /dev/null
+++ b/notebooks/integrations/hugging-face/.nbtest.yml
@@ -0,0 +1,2 @@
+masks:
+- 'Score: [0-9]+\.[0-9][0-9]*'