From 18e2ff2d2d3d464e09fab81b83549167e47e2228 Mon Sep 17 00:00:00 2001 From: Jordan Frazier <122494242+jordanrfrazier@users.noreply.github.com> Date: Mon, 18 Aug 2025 08:25:13 -0400 Subject: [PATCH] fix: fallback to env var correctly when using tweaks (#9422) * Fallback to env var correctly when using tweaks * Add unit test * [autofix.ci] apply automated fixes * refactor: improve docstrings for parameter update tests --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida --- .../langflow/interface/initialize/loading.py | 4 +- src/backend/tests/unit/interface/__init__.py | 1 + .../unit/interface/initialize/__init__.py | 1 + .../unit/interface/initialize/test_loading.py | 225 ++++++++++++++++++ 4 files changed, 230 insertions(+), 1 deletion(-) create mode 100644 src/backend/tests/unit/interface/__init__.py create mode 100644 src/backend/tests/unit/interface/initialize/__init__.py create mode 100644 src/backend/tests/unit/interface/initialize/test_loading.py diff --git a/src/backend/base/langflow/interface/initialize/loading.py b/src/backend/base/langflow/interface/initialize/loading.py index 5a25f8a75..a706a05df 100644 --- a/src/backend/base/langflow/interface/initialize/loading.py +++ b/src/backend/base/langflow/interface/initialize/loading.py @@ -122,7 +122,9 @@ async def update_params_with_load_from_db_fields( try: key = await custom_component.get_variable(name=params[field], field=field, session=session) except ValueError as e: - if any(reason in str(e) for reason in ["User id is not set", "variable not found."]): + if "User id is not set" in str(e): + raise + if "variable not found." in str(e) and not fallback_to_env_vars: raise logger.debug(str(e)) key = None diff --git a/src/backend/tests/unit/interface/__init__.py b/src/backend/tests/unit/interface/__init__.py new file mode 100644 index 000000000..0ad414113 --- /dev/null +++ b/src/backend/tests/unit/interface/__init__.py @@ -0,0 +1 @@ +# Unit tests for interface module diff --git a/src/backend/tests/unit/interface/initialize/__init__.py b/src/backend/tests/unit/interface/initialize/__init__.py new file mode 100644 index 000000000..6ce4cd81e --- /dev/null +++ b/src/backend/tests/unit/interface/initialize/__init__.py @@ -0,0 +1 @@ +# Unit tests for initialize module diff --git a/src/backend/tests/unit/interface/initialize/test_loading.py b/src/backend/tests/unit/interface/initialize/test_loading.py new file mode 100644 index 000000000..7a0860308 --- /dev/null +++ b/src/backend/tests/unit/interface/initialize/test_loading.py @@ -0,0 +1,225 @@ +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langflow.interface.initialize.loading import update_params_with_load_from_db_fields + + +@pytest.mark.asyncio +async def test_update_params_fallback_to_env_when_variable_not_found(): + """Test that when a variable is not found in database and fallback_to_env_vars is True. + + It falls back to environment variables. + """ + # Set up environment variable + os.environ["TEST_API_KEY"] = "test-secret-key-123" + + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(side_effect=ValueError("TEST_API_KEY variable not found.")) + + # Set up params with a field that should load from db + params = {"api_key": "TEST_API_KEY"} + load_from_db_fields = ["api_key"] + + # Call the function with fallback enabled + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + result = await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Should have fallen back to environment variable + assert result["api_key"] == "test-secret-key-123" + + # Clean up + del os.environ["TEST_API_KEY"] + + +@pytest.mark.asyncio +async def test_update_params_raises_when_variable_not_found_and_no_fallback(): + """Test that when a variable is not found and fallback_to_env_vars is False. + + It raises the error. + """ + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(side_effect=ValueError("TEST_API_KEY variable not found.")) + + # Set up params + params = {"api_key": "TEST_API_KEY"} + load_from_db_fields = ["api_key"] + + # Call the function with fallback disabled + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + with pytest.raises(ValueError, match="TEST_API_KEY variable not found"): + await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=False + ) + + +@pytest.mark.asyncio +async def test_update_params_uses_database_variable_when_found(): + """Test that when a variable is found in database, it uses that value. + + It doesn't check environment variables. + """ + # Set up environment variable (should not be used) + os.environ["TEST_API_KEY"] = "env-value" + + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(return_value="db-value") + + # Set up params + params = {"api_key": "TEST_API_KEY"} + load_from_db_fields = ["api_key"] + + # Call the function + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + result = await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Should use database value, not environment value + assert result["api_key"] == "db-value" + + # Clean up + del os.environ["TEST_API_KEY"] + + +@pytest.mark.asyncio +async def test_update_params_sets_none_when_no_env_var_and_fallback_enabled(): + """Test that when variable not found in db and env var doesn't exist. + + The field is set to None. + """ + # Make sure env var doesn't exist + if "NONEXISTENT_KEY" in os.environ: + del os.environ["NONEXISTENT_KEY"] + + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(side_effect=ValueError("NONEXISTENT_KEY variable not found.")) + + # Set up params + params = {"api_key": "NONEXISTENT_KEY"} + load_from_db_fields = ["api_key"] + + # Call the function with fallback enabled + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + result = await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Should be set to None + assert result["api_key"] is None + + +@pytest.mark.asyncio +async def test_update_params_raises_on_user_id_not_set(): + """Test that 'User id is not set' error is always raised regardless of fallback setting.""" + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(side_effect=ValueError("User id is not set")) + + # Set up params + params = {"api_key": "SOME_KEY"} + load_from_db_fields = ["api_key"] + + # Should raise with fallback enabled + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + with pytest.raises(ValueError, match="User id is not set"): + await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Should also raise with fallback disabled + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + with pytest.raises(ValueError, match="User id is not set"): + await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=False + ) + + +@pytest.mark.asyncio +async def test_update_params_skips_empty_fields(): + """Test that empty or None fields in params are skipped.""" + # Create mock custom component + custom_component = MagicMock() + custom_component.get_variable = AsyncMock(return_value="some-value") + + # Set up params with empty and None values + params = {"api_key": "", "another_key": None, "valid_key": "VALID_KEY"} + load_from_db_fields = ["api_key", "another_key", "valid_key"] + + # Call the function + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + result = await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Only valid_key should have been processed + assert result["api_key"] == "" + assert result["another_key"] is None + assert result["valid_key"] == "some-value" + + # get_variable should only be called once for valid_key + custom_component.get_variable.assert_called_once_with( + name="VALID_KEY", field="valid_key", session=mock_session_scope.return_value.__aenter__.return_value + ) + + +@pytest.mark.asyncio +async def test_update_params_handles_multiple_fields(): + """Test that multiple fields are processed correctly with mixed results.""" + # Set up environment variables + os.environ["ENV_KEY"] = "env-value" + + # Create mock custom component + custom_component = MagicMock() + + # Set up different responses for different fields + async def mock_get_variable(name, **_kwargs): + if name == "DB_KEY": + return "db-value" + if name == "ENV_KEY": + error_msg = "ENV_KEY variable not found." + raise ValueError(error_msg) + error_msg = f"{name} variable not found." + raise ValueError(error_msg) + + custom_component.get_variable = AsyncMock(side_effect=mock_get_variable) + + # Set up params + params = {"field1": "DB_KEY", "field2": "ENV_KEY", "field3": "MISSING_KEY"} + load_from_db_fields = ["field1", "field2", "field3"] + + # Call the function + with patch("langflow.interface.initialize.loading.session_scope") as mock_session_scope: + mock_session_scope.return_value.__aenter__.return_value = MagicMock() + + result = await update_params_with_load_from_db_fields( + custom_component, params, load_from_db_fields, fallback_to_env_vars=True + ) + + # Check results + assert result["field1"] == "db-value" # From database + assert result["field2"] == "env-value" # From environment + assert result["field3"] is None # Not found anywhere + + # Clean up + del os.environ["ENV_KEY"]