feat: Add unit tests for run_flow_from_json with fake environment variables (#4015)

* Add tests for run_flow_from_json with fake environment variables

- Implemented test_run_flow_with_fake_env to validate flow execution with a fake .env file.
- Added test_run_flow_with_fake_env_TWEAKS to check flow execution using environment variables loaded from the fake .env file.

* Replace keys in tweaks with their corresponding environment variable values

- Implemented a function to recursively replace keys in the tweaks dictionary with values from the provided environment variables.

* updated to use better way to load test  json file

* [autofix.ci] apply automated fixes

* [autofix.ci] apply automated fixes

* refactor: improve test readability and consistency in load tests

- Renamed variable `TWEAKS` to `tweaks_dict` for clarity and consistency across tests.
- Updated test function names to follow a consistent naming convention.
- Enhanced comments for better understanding of test intentions.
- Minor formatting adjustments to improve code readability.

* feat: add aload_flow_from_json and arun_flow_from_json to module exports

* fix: correct file path handling in aload_flow_from_json function

* fix: improve environment variable handling in aload_flow_from_json function

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>

src/backend/base/langflow/load/__init__.py

@@ -1,11 +1,12 @@
 from .load import aload_flow_from_json, arun_flow_from_json, load_flow_from_json, run_flow_from_json
-from .utils import get_flow, upload_file
+from .utils import get_flow, replace_tweaks_with_env, upload_file
 
 __all__ = [
     "aload_flow_from_json",
     "arun_flow_from_json",
     "get_flow",
     "load_flow_from_json",
+    "replace_tweaks_with_env",
     "run_flow_from_json",
     "upload_file",
 ]

src/backend/base/langflow/load/load.py

@@ -1,13 +1,14 @@
-import asyncio
 import json
+from io import StringIO
 from pathlib import Path
 
 from aiofile import async_open
-from dotenv import load_dotenv
+from dotenv import dotenv_values
 from loguru import logger
 
 from langflow.graph import Graph
 from langflow.graph.schema import RunOutputs
+from langflow.load.utils import replace_tweaks_with_env
 from langflow.logging.logger import configure
 from langflow.processing.process import process_tweaks, run_graph
 from langflow.utils.async_helpers import run_until_complete
@@ -49,14 +50,17 @@ async def aload_flow_from_json(
     configure(log_level=log_level, log_file=log_file_path, disable=disable_logs, async_file=True)
 
     # override env variables with .env file
-    if env_file:
-        await asyncio.to_thread(load_dotenv, env_file, override=True)
+    if env_file and tweaks is not None:
+        async with async_open(Path(env_file), encoding="utf-8") as f:
+            content = await f.read()
+            env_vars = dotenv_values(stream=StringIO(content))
+        tweaks = replace_tweaks_with_env(tweaks=tweaks, env_vars=env_vars)
 
     # Update settings with cache and components path
     await update_settings(cache=cache)
 
     if isinstance(flow, str | Path):
-        async with async_open(Path(flow).name, encoding="utf-8") as f:
+        async with async_open(Path(flow), encoding="utf-8") as f:
             content = await f.read()
             flow_graph = json.loads(content)
     # If input is a dictionary, assume it's a JSON object

src/backend/base/langflow/load/utils.py

@@ -99,5 +99,28 @@ def get_flow(url: str, flow_id: str):
         msg = f"Error retrieving flow: {e}"
         raise UploadError(msg) from e
 
-    msg = f"Error retrieving flow: {response.status_code}"
-    raise UploadError(msg)
+
+def replace_tweaks_with_env(tweaks: dict, env_vars: dict) -> dict:
+    """Replace keys in the tweaks dictionary with their corresponding environment variable values.
+
+    This function recursively traverses the tweaks dictionary and replaces any string keys
+    with their values from the provided environment variables. If a key's value is a dictionary,
+    the function will call itself to handle nested dictionaries.
+
+    Args:
+        tweaks (dict): A dictionary containing keys that may correspond to environment variable names.
+        env_vars (dict): A dictionary of environment variables where keys are variable names
+                         and values are their corresponding values.
+
+    Returns:
+        dict: The updated tweaks dictionary with keys replaced by their environment variable values.
+    """
+    for key, value in tweaks.items():
+        if isinstance(value, dict):
+            # Recursively replace in nested dictionaries
+            tweaks[key] = replace_tweaks_with_env(value, env_vars)
+        elif isinstance(value, str):
+            env_value = env_vars.get(value)  # Get the value from the provided environment variables
+            if env_value is not None:
+                tweaks[key] = env_value
+    return tweaks

src/backend/tests/conftest.py

@@ -103,6 +103,7 @@ def pytest_configure(config):
     pytest.VECTOR_STORE_PATH = data_path / "Vector_store.json"
     pytest.SIMPLE_API_TEST = data_path / "SimpleAPITest.json"
     pytest.MEMORY_CHATBOT_NO_LLM = data_path / "MemoryChatbotNoLLM.json"
+    pytest.ENV_VARIABLE_TEST = data_path / "env_variable_test.json"
     pytest.LOOP_TEST = data_path / "LoopTest.json"
     pytest.CODE_WITH_SYNTAX_ERROR = """
 def get_text():

src/backend/tests/unit/base/load/test_load.py

@@ -1,5 +1,8 @@
 import inspect
+import os
 
+import pytest
+from dotenv import load_dotenv
 from langflow.load import run_flow_from_json
 
 
@@ -26,5 +29,53 @@ def test_run_flow_from_json_params():
     params = func_spec.args + func_spec.kwonlyargs
     assert expected_params.issubset(params), "Not all expected parameters are present in run_flow_from_json"
 
-    # TODO: Add tests by loading a flow and running it need to text with fake llm and check if it returns the
-    #  correct output
+    # TODO: Add tests by loading a flow and running it need to text with fake llm and check if it
+    # returns the correct output
+
+
+@pytest.fixture
+def fake_env_file(tmp_path):
+    # Create a fake .env file
+    env_file = tmp_path / ".env"
+    env_file.write_text("TEST_OP=TESTWORKS")
+    return env_file
+
+
+def test_run_flow_with_fake_env(fake_env_file):
+    # Load the flow from the JSON file
+    # flow_file = Path("src/backend/tests/data/env_variable_test.json")
+    flow_file = pytest.ENV_VARIABLE_TEST
+    tweaks_dict = {"Secret-zIbKs": {"secret_key_input": "TEST_OP"}}
+
+    # Run the flow from JSON, providing the fake env file
+    result = run_flow_from_json(
+        flow=flow_file,
+        input_value="some_input_value",
+        env_file=str(fake_env_file),  # Pass the path of the fake env file
+        tweaks=tweaks_dict,
+    )
+    # Extract and check the output data
+    output_data = result[0].outputs[0].results["message"].data["text"]
+    assert output_data == "TESTWORKS"
+
+
+def test_run_flow_with_fake_env_tweaks(fake_env_file):
+    # Load the flow from the JSON file
+    # flow_file = Path("src/backend/tests/data/env_variable_test.json")
+    flow_file = pytest.ENV_VARIABLE_TEST
+
+    # Load env file and set up tweaks
+
+    load_dotenv(str(fake_env_file))
+    tweaks = {
+        "Secret-zIbKs": {"secret_key_input": os.environ["TEST_OP"]},
+    }
+    # Run the flow from JSON without passing the env_file
+    result = run_flow_from_json(
+        flow=flow_file,
+        input_value="some_input_value",
+        tweaks=tweaks,
+    )
+    # Extract and check the output data
+    output_data = result[0].outputs[0].results["message"].data["text"]
+    assert output_data == "TESTWORKS"