Only pass TILEDB_REST_TOKEN to test_cloud during CI (#327)

Only pass TILEDB_REST_TOKEN to test_cloud and Notebook tests during CI

Author: NikolaosPapailiou

.github/workflows/ci-python.yml

@@ -36,9 +36,10 @@ jobs:
           #pip install -e .
           #pytest
           pip install -r test/ipynb/requirements.txt
+          export TILEDB_REST_TOKEN=$TILEDB_CLOUD_HELPER_VAR
           pytest --nbmake test/ipynb
         env:
-          TILEDB_REST_TOKEN: ${{ secrets.TILEDB_CLOUD_HELPER_VAR }}
+          TILEDB_CLOUD_HELPER_VAR: ${{ secrets.TILEDB_CLOUD_HELPER_VAR }}
         shell: bash -el {0}
       - name: Check tiledb-vector-search version
         run: |

apis/python/test/common.py

@@ -329,10 +329,13 @@ def random_name(name: str) -> str:
 
 
 def check_training_input_vectors(
-    index_uri: str, expected_training_sample_size: int, expected_dimensions: int
+    index_uri: str,
+    expected_training_sample_size: int,
+    expected_dimensions: int,
+    config=None,
 ):
     training_input_vectors_uri = f"{index_uri}/{storage_formats[STORAGE_VERSION]['TRAINING_INPUT_VECTORS_ARRAY_NAME']}"
-    with tiledb.open(training_input_vectors_uri, mode="r") as src_array:
+    with tiledb.open(training_input_vectors_uri, mode="r", config=config) as src_array:
         training_input_vectors = np.transpose(src_array[:, :]["values"])
         assert training_input_vectors.shape[0] == expected_training_sample_size
         assert training_input_vectors.shape[1] == expected_dimensions

apis/python/test/test_cloud.py

@@ -18,9 +18,10 @@ class CloudTests(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        if not os.getenv("TILEDB_REST_TOKEN"):
-            raise ValueError("TILEDB_REST_TOKEN not set")
-        tiledb.cloud.login(token=os.getenv("TILEDB_REST_TOKEN"))
+        token = os.getenv("TILEDB_REST_TOKEN")
+        if os.getenv("TILEDB_CLOUD_HELPER_VAR"):
+            token = os.getenv("TILEDB_CLOUD_HELPER_VAR")
+        tiledb.cloud.login(token=token)
         namespace, storage_path, _ = groups._default_ns_path_cred()
         storage_path = storage_path.replace("//", "/").replace("/", "//", 1)
         rand_name = random_name("vector_search")
@@ -58,7 +59,9 @@ def test_cloud_flat(self):
             mode=Mode.BATCH,
         )
         tiledb_index_uri = groups.info(index_uri).tiledb_uri
-        index = vs.flat_index.FlatIndex(uri=tiledb_index_uri)
+        index = vs.flat_index.FlatIndex(
+            uri=tiledb_index_uri, config=tiledb.cloud.Config().dict()
+        )
 
         _, result_i = index.query(queries, k=k)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
@@ -91,7 +94,10 @@ def test_cloud_ivf_flat(self):
         )
 
         tiledb_index_uri = groups.info(index_uri).tiledb_uri
-        index = vs.ivf_flat_index.IVFFlatIndex(uri=tiledb_index_uri)
+        index = vs.ivf_flat_index.IVFFlatIndex(
+            uri=tiledb_index_uri,
+            config=tiledb.cloud.Config().dict(),
+        )
 
         _, result_i = index.query(queries, k=k, nprobe=nprobe)
         assert accuracy(result_i, gt_i) > MINIMUM_ACCURACY
@@ -183,16 +189,14 @@ def test_cloud_ivf_flat_random_sampling(self):
             training_sample_size=training_sample_size,
             max_sampling_tasks=max_sampling_tasks,
             config=tiledb.cloud.Config().dict(),
-            # TODO Re-enable.
-            #  This is temporarily disabled due to an incompatibility of new ingestion code and previous
-            #  UDF library releases.
-            # mode=Mode.BATCH,
+            mode=Mode.BATCH,
         )
 
         check_training_input_vectors(
             index_uri=index_uri,
             expected_training_sample_size=training_sample_size,
             expected_dimensions=queries.shape[1],
+            config=tiledb.cloud.Config().dict(),
         )
 
         _, result_i = index.query(queries, k=k, nprobe=nprobe)