diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py b/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
index 9a6f03c91..24b1de8a7 100644
--- a/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
+++ b/cognee/tests/cli_tests/cli_unit_tests/test_cli_commands.py
@@ -285,12 +285,11 @@ class TestDeleteCommand:
         assert "all" in actions
         assert "force" in actions
 
-    @patch("cognee.datasets.delete_dataset")
     @patch("cognee.modules.data.methods.get_deletion_counts.get_user")
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
     @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
     def test_execute_delete_dataset_with_confirmation(
-        self, mock_asyncio_run, mock_confirm, get_user_mock, delete_dataset_mock
+        self, mock_asyncio_run, mock_confirm, get_user_mock
     ):
         """Test execute delete dataset with user confirmation"""
         expected_user_id = uuid4()
@@ -298,16 +297,20 @@ class TestDeleteCommand:
 
         get_user_mock.return_value = User(id=expected_user_id)
 
-        command = DeleteCommand()
-        args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
+        mock_cognee = MagicMock()
+        mock_cognee.datasets.delete_dataset = AsyncMock()
 
-        mock_confirm.return_value = True
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = DeleteCommand()
+            args = argparse.Namespace(dataset_id=expected_dataset_id, user_id=expected_user_id)
 
-        command.execute(args)
+            mock_confirm.return_value = True
 
-        delete_dataset_mock.assert_awaited_once_with(
-            dataset_id=expected_dataset_id, user_id=expected_user_id
-        )
+            command.execute(args)
+
+            mock_cognee.datasets.delete_dataset.assert_awaited_once_with(
+                dataset_id=expected_dataset_id, user_id=expected_user_id
+            )
 
     @patch("cognee.cli.commands.delete_command.get_deletion_counts")
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
diff --git a/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py b/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py
index 178c6aa7f..e43233a4b 100644
--- a/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py
+++ b/cognee/tests/cli_tests/cli_unit_tests/test_cli_edge_cases.py
@@ -382,25 +382,25 @@ class TestDeleteCommandEdgeCases:
 
     @patch("cognee.cli.commands.delete_command.asyncio.run", side_effect=_mock_run)
     @patch("cognee.cli.commands.delete_command.fmt.confirm")
-    @patch("cognee.datasets.delete_all")
-    def test_delete_all_with_user_id(self, delete_all_mock, fmt_confirm_mock, async_run_mock):
+    def test_delete_all_with_user_id(self, fmt_confirm_mock, async_run_mock):
         """Test delete command with both --all and --user-id"""
         fmt_confirm_mock.return_value = True
-        delete_all_mock.return_value = None
 
         expected_user_id = uuid4()
 
-        command = DeleteCommand()
-        args = argparse.Namespace(
-            dataset_name=None,
-            user_id=expected_user_id,
-            all=True,
-            force=False,
-        )
+        mock_cognee = MagicMock()
+        mock_cognee.datasets.delete_all = AsyncMock()
 
-        command.execute(args)
+        with patch.dict(sys.modules, {"cognee": mock_cognee}):
+            command = DeleteCommand()
+            args = argparse.Namespace(
+                user_id=expected_user_id,
+                all=True,
+            )
 
-        delete_all_mock.assert_called_once_with(user_id=expected_user_id)
+            command.execute(args)
+
+            mock_cognee.datasets.delete_all.assert_called_once_with(user_id=expected_user_id)
 
     @patch("cognee.cli.commands.delete_command.get_deletion_counts")
     @patch("cognee.cli.commands.delete_command.fmt.confirm")