From 56731ad08cb297f87f5dc90ebf5590ecef6331e4 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Fri, 8 Mar 2024 16:43:21 -0300
Subject: [PATCH] Add tests for directory component and multithreading

---
 tests/test_data_components.py | 66 +++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/tests/test_data_components.py b/tests/test_data_components.py
index 6aa9a4768..3d8838514 100644
--- a/tests/test_data_components.py
+++ b/tests/test_data_components.py
@@ -1,3 +1,7 @@
+import os
+from pathlib import Path
+from unittest.mock import Mock, patch
+
 import httpx
 import pytest
 import respx
@@ -91,3 +95,65 @@ async def test_build_with_multiple_urls(api_request):
 
     # Assertions
     assert len(results) == len(urls)
+
+
+@patch("langflow.components.data.Directory.parallel_load_records")
+@patch("langflow.components.data.Directory.retrieve_file_paths")
+@patch("langflow.components.data.DirectoryComponent.resolve_path")
+def test_directory_component_build_with_multithreading(
+    mock_resolve_path, mock_retrieve_file_paths, mock_parallel_load_records
+):
+    # Arrange
+    directory_component = data.DirectoryComponent()
+    path = os.path.dirname(os.path.abspath(__file__))
+    types = ["py"]
+    depth = 1
+    max_concurrency = 2
+    load_hidden = False
+    recursive = True
+    silent_errors = False
+    use_multithreading = True
+
+    mock_resolve_path.return_value = path
+    mock_retrieve_file_paths.return_value = [
+        os.path.join(path, file) for file in os.listdir(path) if file.endswith(".py")
+    ]
+    mock_parallel_load_records.return_value = [Mock()]
+
+    # Act
+    result = directory_component.build(
+        path,
+        types,
+        depth,
+        max_concurrency,
+        load_hidden,
+        recursive,
+        silent_errors,
+        use_multithreading,
+    )
+
+    # Assert
+    mock_resolve_path.assert_called_once_with(path)
+    mock_retrieve_file_paths.assert_called_once_with(
+        path, types, load_hidden, recursive, depth
+    )
+    mock_parallel_load_records.assert_called_once_with(
+        mock_retrieve_file_paths.return_value, silent_errors, max_concurrency
+    )
+
+
+def test_directory_without_mocks():
+    directory_component = data.DirectoryComponent()
+    from langflow.initial_setup import setup
+    from langflow.initial_setup.setup import load_starter_projects
+
+    projects = load_starter_projects()
+    # the setup module has a folder where the projects are stored
+    # the contents of that folder are in the projects variable
+    # the directory component can be used to load the projects
+    # and we can validate if the contents are the same as the projects variable
+    setup_path = Path(setup.__file__).parent / "starter_projects"
+    result = directory_component.build(
+        str(setup_path), types=["json"], use_multithreading=False
+    )
+    assert len(result) == len(projects)