fix: Rewrite environment tests to properly mock dependencies

src/strands_tools/environment.py

@@ -71,8 +71,7 @@
 from rich.text import Text
 from strands.types.tools import ToolResult, ToolResultContent, ToolUse
 
-from strands_tools.utils import console_util
-from strands_tools.utils.user_input import get_user_input
+from strands_tools.utils import console_util, user_input
 
 TOOL_SPEC = {
     "name": "environment",
@@ -584,7 +583,7 @@ def environment(tool: ToolUse, **kwargs: Any) -> ToolResult:
                 )
 
             # Ask for confirmation
-            confirm = get_user_input(
+            confirm = user_input.get_user_input(
                 "\n<yellow><bold>Do you want to proceed with setting this environment variable?</bold> [y/*]</yellow>"
             )
             # For tests, 'y' should be recognized even with extra spaces or newlines
@@ -706,7 +705,7 @@ def environment(tool: ToolUse, **kwargs: Any) -> ToolResult:
                 )
 
                 # Ask for confirmation
-                confirm = get_user_input(
+                confirm = user_input.get_user_input(
                     "\n<red><bold>Do you want to proceed with deleting this environment variable?</bold> [y/*]</red>"
                 )
                 # For tests, 'y' should be recognized even with extra spaces or newlines

tests/test_environment.py

@@ -3,28 +3,33 @@
 """
 
 import os
+from unittest import mock
 
 import pytest
 from strands import Agent
 from strands_tools import environment
+from strands_tools.utils import user_input
 
 
 @pytest.fixture
 def agent():
     """Create an agent with the environment tool loaded."""
-    return Agent(tools=[environment])
+    return Agent(tools=[environment], load_tools_from_directory=False)
 
 
-@pytest.fixture
-def test_env_var():
-    """Create and clean up a test environment variable."""
-    var_name = "TEST_ENV_VAR"
-    var_value = "test_value"
-    os.environ[var_name] = var_value
-    yield var_name, var_value
-    # Clean up: remove test variable if it exists
-    if var_name in os.environ:
-        del os.environ[var_name]
+@pytest.fixture(autouse=True)
+def get_user_input():
+    with mock.patch.object(user_input, "get_user_input") as mocked_user_input:
+        # By default all tests will return deny
+        mocked_user_input.return_value = "n"
+        yield mocked_user_input
+
+
+@pytest.fixture(autouse=True)
+def os_environment():
+    mock_env = {}
+    with mock.patch.object(os, "environ", mock_env):
+        yield mock_env
 
 
 def extract_result_text(result):
@@ -42,18 +47,24 @@ def test_direct_list_action(agent):
     assert len(extract_result_text(result)) > 0
 
 
-def test_direct_list_with_prefix(agent, test_env_var):
+def test_direct_list_with_prefix(agent, os_environment):
     """Test listing environment variables with a specific prefix."""
-    var_name, _ = test_env_var
+    var_name = "TEST_ENV_VAR"
+    var_value = "test_value"
+    os_environment[var_name] = var_value
+
     result = agent.tool.environment(action="list", prefix=var_name[:4])
     assert result["status"] == "success"
     # Verify our test variable is in the result
     assert var_name in extract_result_text(result)
 
 
-def test_direct_get_existing_var(agent, test_env_var):
+def test_direct_get_existing_var(agent, os_environment):
     """Test getting an existing environment variable."""
-    var_name, var_value = test_env_var
+    var_name = "TEST_ENV_VAR"
+    var_value = "test_value"
+    os_environment[var_name] = var_value
+
     result = agent.tool.environment(action="get", name=var_name)
     assert result["status"] == "success"
     assert var_name in extract_result_text(result)
@@ -67,10 +78,9 @@ def test_direct_get_nonexistent_var(agent):
     assert "not found" in extract_result_text(result)
 
 
-def test_direct_set_protected_var(agent, monkeypatch):
+def test_direct_set_protected_var(agent, os_environment):
     """Test attempting to set a protected environment variable."""
-    # Mock get_user_input to always return 'y' for confirmation
-    monkeypatch.setattr("strands_tools.utils.user_input.get_user_input", lambda _: "y")
+    os_environment["PATH"] = "/original/path"
 
     # Try to modify PATH which is in PROTECTED_VARS
     result = agent.tool.environment(action="set", name="PATH", value="/bad/path")
@@ -79,103 +89,91 @@ def test_direct_set_protected_var(agent, monkeypatch):
     assert os.environ["PATH"] != "/bad/path"
 
 
-@pytest.mark.skipif(os.name == "nt", reason="Need to mock console output for Windows (see #17)")
-def test_direct_set_var_cancelled(agent, monkeypatch):
-    """Test cancelling setting an environment variable."""
-    # Mock get_user_input to return 'n' to cancel
-    monkeypatch.setattr("strands_tools.utils.user_input.get_user_input", lambda _: "n")
+def test_direct_set_var_allowed(get_user_input, agent):
+    """Test attempting to set a protected environment variable and allowing it."""
+    get_user_input.return_value = "y"
 
     var_name = "CANCELLED_VAR"
     var_value = "cancelled_value"
 
-    # Clean up in case the variable exists
-    if var_name in os.environ:
-        del os.environ[var_name]
+    result = agent.tool.environment(action="set", name=var_name, value=var_value)
+    assert result["status"] == "success"
+    assert var_name in os.environ
+    assert get_user_input.call_count == 1
+
+
+def test_direct_set_var_cancelled(agent):
+    var_name = "CANCELLED_VAR"
+    var_value = "cancelled_value"
 
-    try:
-        result = agent.tool.environment(action="set", name=var_name, value=var_value)
-        assert result["status"] == "error"
-        assert "cancelled" in extract_result_text(result).lower()
-        # Verify variable was not set
-        assert var_name not in os.environ
-    finally:
-        # Clean up
-        if var_name in os.environ:
-            del os.environ[var_name]
+    result = agent.tool.environment(action="set", name=var_name, value=var_value)
+    assert result["status"] == "error"
+    assert "cancelled" in extract_result_text(result).lower()
+    # Verify variable was not set
+    assert var_name not in os.environ
 
 
-def test_direct_delete_nonexistent_var(agent, monkeypatch):
+def test_direct_delete_nonexistent_var(agent):
     """Test attempting to delete a non-existent variable."""
-    # Mock get_user_input to always return 'y' for confirmation
-    monkeypatch.setattr("strands_tools.utils.user_input.get_user_input", lambda _: "y")
-
     var_name = "NONEXISTENT_VAR_FOR_DELETE_TEST"
 
-    # Make sure the variable doesn't exist
-    if var_name in os.environ:
-        del os.environ[var_name]
-
     result = agent.tool.environment(action="delete", name=var_name)
     assert result["status"] == "error"
     assert "not found" in extract_result_text(result).lower()
 
 
-def test_direct_delete_protected_var(agent, monkeypatch):
+def test_direct_delete_protected_var(agent, os_environment):
     """Test attempting to delete a protected environment variable."""
-    # Mock get_user_input to always return 'y' for confirmation
-    monkeypatch.setattr("strands_tools.utils.user_input.get_user_input", lambda _: "y")
-
     # Try to delete PATH which is in PROTECTED_VARS
-    original_path = os.environ.get("PATH", "")
-    try:
-        result = agent.tool.environment(action="delete", name="PATH")
-        assert result["status"] == "error"
-        # Verify PATH still exists
-        assert "PATH" in os.environ
-    finally:
-        # Restore PATH if somehow it got deleted
-        if "PATH" not in os.environ:
-            os.environ["PATH"] = original_path
-
-
-@pytest.mark.skipif(os.name == "nt", reason="Need to mock console output for Windows (see #17)")
-def test_direct_delete_var_cancelled(agent, monkeypatch):
+    unchanging_value = "/original/path"
+    os_environment["PATH"] = unchanging_value
+
+    result = agent.tool.environment(action="delete", name="PATH")
+    assert result["status"] == "error"
+    # Verify PATH still exists
+    assert os_environment["PATH"] == unchanging_value
+
+
+def test_direct_delete_var_cancelled(agent, os_environment):
     """Test cancelling deletion of an environment variable."""
-    # Mock get_user_input to return 'n' to cancel
-    monkeypatch.setattr("strands_tools.utils.user_input.get_user_input", lambda _: "n")
+    var_name = "CANCEL_DELETE_VAR"
+    var_value = "cancel_delete_value"
+
+    # Set up the variable
+    os_environment[var_name] = var_value
 
-    # Ensure DEV mode is disabled to force confirmation
-    current_dev = os.environ.get("DEV", None)
-    if current_dev:
-        os.environ.pop("DEV")
+    result = agent.tool.environment(action="delete", name=var_name)
+    assert result["status"] == "error"
+    assert "cancelled" in extract_result_text(result).lower()
+    # Verify variable still exists
+    assert var_name in os.environ
+    assert os_environment[var_name] == var_value
+
+
+def test_direct_delete_var_allowed(agent, get_user_input, os_environment):
+    """Test allowing deletion of an environment variable."""
+    get_user_input.return_value = "y"
 
     var_name = "CANCEL_DELETE_VAR"
     var_value = "cancel_delete_value"
 
     # Set up the variable
-    os.environ[var_name] = var_value
-
-    try:
-        result = agent.tool.environment(action="delete", name=var_name)
-        assert result["status"] == "error"
-        assert "cancelled" in extract_result_text(result).lower()
-        # Verify variable still exists
-        assert var_name in os.environ
-        assert os.environ[var_name] == var_value
-    finally:
-        # Clean up
-        if var_name in os.environ:
-            del os.environ[var_name]
-        # Restore DEV mode if it was set
-        if current_dev:
-            os.environ["DEV"] = current_dev
-        if var_name in os.environ:
-            del os.environ[var_name]
-
-
-def test_direct_validate_existing_var(agent, test_env_var):
+    os_environment[var_name] = var_value
+
+    result = agent.tool.environment(action="delete", name=var_name)
+    assert result["status"] == "success"
+    assert "deleted environment variable" in extract_result_text(result).lower()
+    assert os_environment.get(var_name) is None
+
+
+def test_direct_validate_existing_var(agent, os_environment):
     """Test validating an existing environment variable."""
-    var_name, _ = test_env_var
+    var_name = "CANCEL_DELETE_VAR"
+    var_value = "cancel_delete_value"
+
+    # Set up the variable
+    os_environment[var_name] = var_value
+
     result = agent.tool.environment(action="validate", name=var_name)
     assert result["status"] == "success"
     assert "valid" in extract_result_text(result).lower()
@@ -203,76 +201,52 @@ def test_direct_missing_parameters(agent):
     assert result["status"] == "error"
 
 
-def test_environment_dev_mode_delete(agent):
+def test_environment_dev_mode_delete(agent, os_environment):
     """Test the environment tool in DEV mode with delete action."""
     # Set DEV mode
-    original_dev = os.environ.get("DEV")
-    os.environ["DEV"] = "true"
+    os_environment["DEV"] = "true"
 
     var_name = "DEV_MODE_DELETE_VAR"
     var_value = "dev_mode_delete_value"
 
-    try:
-        # Set up the variable
-        os.environ[var_name] = var_value
-
-        result = agent.tool.environment(action="delete", name=var_name)
-        assert result["status"] == "success"
-        assert var_name not in os.environ
-    finally:
-        # Clean up
-        if var_name in os.environ:
-            del os.environ[var_name]
+    # Set up the variable
+    os_environment[var_name] = var_value
 
-        # Restore original DEV value
-        if original_dev is None:
-            if "DEV" in os.environ:
-                del os.environ["DEV"]
-        else:
-            os.environ["DEV"] = original_dev
+    result = agent.tool.environment(action="delete", name=var_name)
+    assert result["status"] == "success"
+    assert var_name not in os_environment
 
 
-def test_environment_dev_mode_protected_var(agent, monkeypatch):
+def test_environment_dev_mode_protected_var(agent, os_environment):
     """Test that protected variables are still protected in DEV mode."""
     # Set DEV mode
-    original_dev = os.environ.get("DEV")
-    os.environ["DEV"] = "true"
-
-    try:
-        # Try to modify PATH which is protected
-        result = agent.tool.environment(action="set", name="PATH", value="/bad/path")
-        assert result["status"] == "error"
-        # Verify PATH was not changed
-        assert os.environ["PATH"] != "/bad/path"
-    finally:
-        # Restore original DEV value
-        if original_dev is None:
-            if "DEV" in os.environ:
-                del os.environ["DEV"]
-        else:
-            os.environ["DEV"] = original_dev
-
-
-def test_environment_masked_values(agent, test_env_var):
+    os_environment["DEV"] = True
+
+    unchanging_value = "/original/path"
+    os_environment["PATH"] = unchanging_value
+
+    # Try to modify PATH which is protected
+    result = agent.tool.environment(action="set", name="PATH", value="/bad/path")
+    assert result["status"] == "error"
+    # Verify PATH was not changed
+    assert os_environment != unchanging_value
+
+
+def test_environment_masked_values(agent, os_environment):
     """Test that sensitive values are masked in output."""
     # Create a sensitive looking variable
     sensitive_name = "TEST_TOKEN_SECRET"
     sensitive_value = "abcd1234efgh5678"
-    os.environ[sensitive_name] = sensitive_value
-
-    try:
-        # Test with masking enabled (default)
-        result = agent.tool.environment(action="get", name=sensitive_name)
-        assert result["status"] == "success"
-        # The full value should not appear in the output
-        assert sensitive_value not in extract_result_text(result)
-
-        # Test with masking disabled
-        result = agent.tool.environment(action="get", name=sensitive_name, masked=False)
-        assert result["status"] == "success"
-        # Now the full value should appear
-        assert sensitive_value in extract_result_text(result)
-    finally:
-        # Clean up
-        if sensitive_name in os.environ:
-            del os.environ[sensitive_name]
+    os_environment[sensitive_name] = sensitive_value
+
+    # Test with masking enabled (default)
+    result = agent.tool.environment(action="get", name=sensitive_name)
+    assert result["status"] == "success"
+    # The full value should not appear in the output
+    assert sensitive_value not in extract_result_text(result)
+
+    # Test with masking disabled
+    result = agent.tool.environment(action="get", name=sensitive_name, masked=False)
+    assert result["status"] == "success"
+    # Now the full value should appear
+    assert sensitive_value in extract_result_text(result)