Fix taskquery tests

ORNL · Dec 12, 2024 · 12927de · 12927de
1 parent 5c06e4a
commit 12927de
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 13 deletions.
diff --git a/.github/workflows/run-tests.yml b/.github/workflows/run-tests.yml
@@ -44,7 +44,7 @@ jobs:
         run: |
           python examples/simple_instrumented_script.py | tee output.log
           cat output.log
-          grep -q "ERROR" output.log && exit 1
+          grep -iq "error" output.log && exit 1
           rm output.log
  
 
@@ -54,7 +54,7 @@ jobs:
           pip install .[dask]
           python examples/dask_example.py | tee output.log
           cat output.log
-          grep -q "ERROR" output.log && exit 1
+          grep -iq "error" output.log && exit 1
           rm output.log          
 
       - name: Install MLFlow dependencies alone and run a simple MLFlow test
@@ -63,7 +63,7 @@ jobs:
           pip install .[mlflow]
           python examples/mlflow_example.py | tee output.log
           cat output.log
-          grep -q "ERROR" output.log && exit 1
+          grep -iq "error" output.log && exit 1
           rm output.log    
 
       - name: Install Tensorboard dependencies alone and run a simple Tensorboard test
@@ -72,7 +72,7 @@ jobs:
           pip install .[tensorboard]
           python examples/tensorboard_example.py  | tee output.log
           cat output.log
-          grep -q "ERROR" output.log && exit 1
+          grep -iq "error" output.log && exit 1
           rm output.log 
 
       - name: Install all dependencies

diff --git a/examples/dask_example.py b/examples/dask_example.py
@@ -65,6 +65,6 @@ def sum_list(values):
     print(all_tasks)
     print("\n\n")
     print("Getting workflow info:")
-    wf_info = Flowcept.db.query(filter={"workflow_id": wf_id}, type="workflow")[0]
+    wf_info = Flowcept.db.query(filter={"workflow_id": wf_id}, collection="workflows")[0]
     assert wf_info["workflow_id"] == wf_id
     print(wf_info)
diff --git a/src/flowcept/flowcept_api/task_query_api.py b/src/flowcept/flowcept_api/task_query_api.py
@@ -12,6 +12,7 @@
 
 from bson.objectid import ObjectId
 
+from flowcept import Flowcept
 from flowcept.analytics.analytics_utils import (
     clean_dataframe as clean_df,
     analyze_correlations_between,
@@ -23,7 +24,6 @@
     to_datetime,
     calculate_telemetry_diff_for_docs,
 )
-from flowcept.flowcept_api.db_api import DBAPI
 from flowcept.configs import WEBSERVER_HOST, WEBSERVER_PORT, ANALYTICS
 from flowcept.flowcept_webserver.app import BASE_ROUTE
 from flowcept.flowcept_webserver.resources.query_rsrc import TaskQuery
@@ -135,7 +135,7 @@ def query(
                 raise Exception(r.text)
 
         else:
-            db_api = DBAPI()
+            db_api = Flowcept.db
             docs = db_api.task_query(
                 filter,
                 projection,
@@ -151,7 +151,7 @@ def query(
 
     def get_subworkflows_tasks_from_a_parent_workflow(self, parent_workflow_id: str) -> List[Dict]:
         """Get subworkflows."""
-        db_api = DBAPI()
+        db_api = Flowcept.db
         sub_wfs = db_api.workflow_query({"parent_workflow_id": parent_workflow_id})
         if not sub_wfs:
             return None

diff --git a/tests/api/task_query_api_test.py b/tests/api/task_query_api_test.py
@@ -146,20 +146,21 @@ def __init__(self, *args, **kwargs):
         super(TaskQueryAPITest, self).__init__(*args, **kwargs)
         self.logger = FlowceptLogger()
         self.api = TaskQueryAPI()
-        self.db_dao = DocumentDBDAO.get_instance(create_indices=False)
 
     def gen_n_get_task_ids(self, generation_function, size=1, generation_args={}):
         docs, task_ids = generation_function(size=size, **generation_args)
 
-        init_db_count = self.db_dao.count_tasks()
-        self.db_dao.insert_and_update_many_tasks(docs, "task_id")
+        dao = DocumentDBDAO.get_instance(create_indices=False)
+        init_db_count = dao.count_tasks()
+        dao.insert_and_update_many_tasks(docs, "task_id")
 
         task_ids_filter = {"task_id": {"$in": task_ids}}
         return task_ids_filter, task_ids, init_db_count
 
     def delete_task_ids_and_assert(self, task_ids, init_db_count):
-        self.db_dao.delete_task_keys("task_id", task_ids)
-        final_db_count = self.db_dao.count_tasks()
+        dao = DocumentDBDAO.get_instance(create_indices=False)
+        dao.delete_task_keys("task_id", task_ids)
+        final_db_count = dao.count_tasks()
         assert init_db_count == final_db_count
 
     def test_webserver_query(self):