adding all tasks test support

NVIDIA · Nov 27, 2024 · b7e40e5 · b7e40e5
1 parent cffb1e4
commit b7e40e5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/client/src/nv_ingest_client/client/interface.py b/client/src/nv_ingest_client/client/interface.py
@@ -302,7 +302,8 @@ def all_tasks(self) -> "Ingestor":
             .dedup() \
             .filter() \
             .split() \
-            .embed()
+            .embed() \
+            .store_embed()
         # .store() \
         # .vdb_upload()
         # fmt: on

diff --git a/tests/nv_ingest_client/client/test_interface.py b/tests/nv_ingest_client/client/test_interface.py
@@ -330,7 +330,7 @@ def test_files_with_remote_files(ingestor_without_doc):
 def test_all_tasks_adds_default_tasks(ingestor):
     ingestor.all_tasks()
 
-    task_classes = {ExtractTask, DedupTask, FilterTask, SplitTask, EmbedTask}
+    task_classes = {ExtractTask, DedupTask, FilterTask, SplitTask, EmbedTask, StoreEmbedTask}
     added_tasks = {
         type(task) for job_specs in ingestor._job_specs._file_type_to_job_spec.values() for task in job_specs[0]._tasks
     }