add new test

Author: jtwinrelevanceai

ai_transform/engine/abstract_engine.py

@@ -1,5 +1,4 @@
 import time
-import logging
 import warnings
 
 from json import JSONDecodeError
@@ -8,7 +7,7 @@
 
 from tqdm.auto import tqdm
 
-from ai_transform.logger import format_logging_info, ic
+from ai_transform.logger import ic
 from ai_transform.types import Filter
 from ai_transform.dataset.dataset import Dataset
 from ai_transform.operator.abstract_operator import AbstractOperator

tests/conftest.py

@@ -63,6 +63,20 @@ def full_dataset(test_client: Client) -> Dataset:
     test_client.delete_dataset(dataset_id)
 
 
+@pytest.fixture(scope="class")
+def partial_dataset(test_client: Client) -> Dataset:
+    salt = "".join(random.choices(string.ascii_lowercase, k=10))
+    dataset_id = f"_sample_dataset_{salt}"
+    dataset = test_client.Dataset(dataset_id, expire=True)
+    documents = mock_documents(1000)
+    for document in documents:
+        for field in random.choices(document.keys(), k=min(len(document), 5)):
+            document.pop(field)
+    dataset.insert_documents(documents)
+    yield dataset
+    test_client.delete_dataset(dataset_id)
+
+
 @pytest.fixture(scope="class")
 def mixed_dataset(test_client: Client) -> Dataset:
     salt = "".join(random.choices(string.ascii_lowercase, k=10))

tests/core/test_engine/test_stable_engine.py

@@ -3,18 +3,24 @@
 from ai_transform.engine.small_batch_stable_engine import SmallBatchStableEngine
 
 from ai_transform.operator.abstract_operator import AbstractOperator
-from ai_transform.workflow.abstract_workflow import AbstractWorkflow
+from ai_transform.workflow.abstract_workflow import Workflow
 
 
 class TestStableEngine:
     def test_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = StableEngine(full_dataset, test_operator, worker_number=0)
-        workflow = AbstractWorkflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
         workflow.run()
         assert engine.success_ratio == 1
 
     def test_small_batch_stable_engine(self, full_dataset: Dataset, test_operator: AbstractOperator):
         engine = SmallBatchStableEngine(full_dataset, test_operator)
-        workflow = AbstractWorkflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
+        workflow.run()
+        assert engine.success_ratio == 1
+
+    def test_stable_engine_filters(self, partial_dataset: Dataset, test_operator: AbstractOperator):
+        engine = StableEngine(partial_dataset, test_operator, select_fields=["sample_1_label"])
+        workflow = Workflow(name="workflow_test123", engine=engine, job_id="test_job123")
         workflow.run()
         assert engine.success_ratio == 1