diff --git a/cognee/modules/tasks/create_task_status_table.py b/cognee/modules/tasks/create_task_status_table.py
index ec9e045aa..26beceb1c 100644
--- a/cognee/modules/tasks/create_task_status_table.py
+++ b/cognee/modules/tasks/create_task_status_table.py
@@ -1,7 +1,7 @@
-from cognee.infrastructure.databases.relational.config import get_relationaldb_config
-config = get_relationaldb_config()
+from cognee.infrastructure.InfrastructureConfig import infrastructure_config
+
 def create_task_status_table():
-    db_engine = config.db_engine
+    db_engine = infrastructure_config.get_config("database_engine")
 
     db_engine.create_table("cognee_task_status", [
         dict(name = "data_id", type = "STRING"),
diff --git a/cognee/tests/test_library.py b/cognee/tests/test_library.py
index cc02e533a..eaac864bc 100755
--- a/cognee/tests/test_library.py
+++ b/cognee/tests/test_library.py
@@ -25,7 +25,6 @@ async def  main():
     dataset_name = "cs_explanations"
 
     explanation_file_path = os.path.join(pathlib.Path(__file__).parent, "test_data/Natural_language_processing.txt")
-    #
     await cognee.add([explanation_file_path], dataset_name)
     
     # dataset_name = "short_stories"