diff --git a/holmes/config.py b/holmes/config.py index 42e7ac7a66..9d548855cd 100644 --- a/holmes/config.py +++ b/holmes/config.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union import yaml # type: ignore -from pydantic import BaseModel, ConfigDict, FilePath, SecretStr +from pydantic import BaseModel, ConfigDict, FilePath, SecretStr, field_validator from holmes.common.env_vars import ROBUSTA_AI, ROBUSTA_API_ENDPOINT, ROBUSTA_CONFIG_PATH from holmes.core.tools_utils.tool_executor import ToolExecutor @@ -112,6 +112,28 @@ class Config(RobustaBaseConfig): # custom_toolsets_from_cli is passed from CLI option `--custom-toolsets` as 'experimental' custom toolsets. # The status of toolset here won't be cached, so the toolset from cli will always be loaded when specified in the CLI. custom_toolsets_from_cli: Optional[List[FilePath]] = None + + # Optional filter for builtin toolsets - specifies which builtin toolsets should be loaded + allowed_builtin_toolsets: Optional[List[str]] = None + + @field_validator("allowed_builtin_toolsets") + @classmethod + def _validate_allowed_builtin_toolsets( + cls, v: Optional[List[str]] + ) -> Optional[List[str]]: + """Validate allowed_builtin_toolsets field.""" + if v is None: + return v + + # Filter out empty strings and whitespace-only strings + if isinstance(v, list): + filtered = [name.strip() for name in v if name and name.strip()] + # Important: preserve empty list as empty list (different from None) + # None = no filtering, [] = filter to nothing + return filtered + + return v + should_try_robusta_ai: bool = False # if True, we will try to load the Robusta AI model, in cli we aren't trying to load it. toolsets: Optional[dict[str, dict[str, Any]]] = None @@ -125,6 +147,7 @@ class Config(RobustaBaseConfig): def toolset_manager(self) -> ToolsetManager: if not self._toolset_manager: self._toolset_manager = ToolsetManager( + config=self, # Pass self as config parameter toolsets=self.toolsets, mcp_servers=self.mcp_servers, custom_toolsets=self.custom_toolsets, @@ -182,6 +205,16 @@ def load_from_file(cls, config_file: Optional[Path], **kwargs) -> "Config": cli_options = {k: v for k, v in kwargs.items() if v is not None and v != []} + # Parse CLI option for allowed_builtin_toolsets if provided + if "allowed_builtin_toolsets" in cli_options and isinstance( + cli_options["allowed_builtin_toolsets"], str + ): + cli_options["allowed_builtin_toolsets"] = [ + name.strip() + for name in cli_options["allowed_builtin_toolsets"].split(",") + if name.strip() + ] + if config_from_file is None: result = cls(**cli_options) else: diff --git a/holmes/core/toolset_manager.py b/holmes/core/toolset_manager.py index 7177cdac4b..c891e1a1ae 100644 --- a/holmes/core/toolset_manager.py +++ b/holmes/core/toolset_manager.py @@ -25,12 +25,14 @@ class ToolsetManager: def __init__( self, + config: Optional[Any] = None, # Config instance toolsets: Optional[dict[str, dict[str, Any]]] = None, mcp_servers: Optional[dict[str, dict[str, Any]]] = None, custom_toolsets: Optional[List[FilePath]] = None, custom_toolsets_from_cli: Optional[List[FilePath]] = None, toolset_status_location: Optional[FilePath] = None, ): + self._config = config # Store config instance self.toolsets = toolsets self.toolsets = toolsets or {} if mcp_servers is not None: @@ -65,6 +67,12 @@ def server_tool_tags(self) -> List[ToolsetTag]: """ return [ToolsetTag.CORE, ToolsetTag.CLUSTER] + def _get_allowed_builtin_toolsets(self) -> Optional[List[str]]: + """Get allowed builtin toolsets from config.""" + if self._config is None: + return None + return getattr(self._config, "allowed_builtin_toolsets", None) + def _list_all_toolsets( self, dal: Optional[SupabaseDal] = None, @@ -81,7 +89,9 @@ def _list_all_toolsets( 3. custom toolset from config can override both built-in and add new custom toolsets # for backward compatibility """ # Load built-in toolsets - builtin_toolsets = load_builtin_toolsets(dal) + builtin_toolsets = load_builtin_toolsets( + dal, allowed_builtin_toolsets=self._get_allowed_builtin_toolsets() + ) toolsets_by_name: dict[str, Toolset] = { toolset.name: toolset for toolset in builtin_toolsets } diff --git a/holmes/main.py b/holmes/main.py index 9f8adecf25..303aa89e40 100644 --- a/holmes/main.py +++ b/holmes/main.py @@ -93,6 +93,11 @@ "-r", help="Path to a custom runbooks (can specify -r multiple times to add multiple runbooks)", ) +opt_allowed_builtin_toolsets: Optional[str] = typer.Option( + None, + "--allowed-builtin-toolsets", + help="Comma-separated list of builtin toolsets to allow (e.g., 'kubernetes/core,prometheus/core')", +) opt_max_steps: Optional[int] = typer.Option( 10, "--max-steps", @@ -174,6 +179,7 @@ def ask( model: Optional[str] = opt_model, config_file: Optional[Path] = opt_config_file, custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, # semi-common options @@ -240,6 +246,7 @@ def ask( model=model, max_steps=max_steps, custom_toolsets_from_cli=custom_toolsets, + allowed_builtin_toolsets=allowed_builtin_toolsets, slack_token=slack_token, slack_channel=slack_channel, ) @@ -371,6 +378,7 @@ def alertmanager( config_file: Optional[Path] = opt_config_file, # type: ignore custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, custom_runbooks: Optional[List[Path]] = opt_custom_runbooks, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, # advanced options for this command @@ -402,6 +410,7 @@ def alertmanager( slack_channel=slack_channel, custom_toolsets_from_cli=custom_toolsets, custom_runbooks=custom_runbooks, + allowed_builtin_toolsets=allowed_builtin_toolsets, ) ai = config.create_console_issue_investigator() # type: ignore @@ -508,6 +517,7 @@ def jira( config_file: Optional[Path] = opt_config_file, # type: ignore custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, custom_runbooks: Optional[List[Path]] = opt_custom_runbooks, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, json_output_file: Optional[str] = opt_json_output_file, @@ -532,6 +542,7 @@ def jira( jira_query=jira_query, custom_toolsets_from_cli=custom_toolsets, custom_runbooks=custom_runbooks, + allowed_builtin_toolsets=allowed_builtin_toolsets, ) ai = config.create_console_issue_investigator() # type: ignore source = config.create_jira_source() @@ -699,6 +710,7 @@ def github( config_file: Optional[Path] = opt_config_file, # type: ignore custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, custom_runbooks: Optional[List[Path]] = opt_custom_runbooks, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, # advanced options for this command @@ -723,6 +735,7 @@ def github( github_query=github_query, custom_toolsets_from_cli=custom_toolsets, custom_runbooks=custom_runbooks, + allowed_builtin_toolsets=allowed_builtin_toolsets, ) ai = config.create_console_issue_investigator() source = config.create_github_source() @@ -784,6 +797,7 @@ def pagerduty( config_file: Optional[Path] = opt_config_file, # type: ignore custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, custom_runbooks: Optional[List[Path]] = opt_custom_runbooks, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, json_output_file: Optional[str] = opt_json_output_file, @@ -807,6 +821,7 @@ def pagerduty( pagerduty_incident_key=pagerduty_incident_key, custom_toolsets_from_cli=custom_toolsets, custom_runbooks=custom_runbooks, + allowed_builtin_toolsets=allowed_builtin_toolsets, ) ai = config.create_console_issue_investigator() source = config.create_pagerduty_source() @@ -870,6 +885,7 @@ def opsgenie( config_file: Optional[Path] = opt_config_file, # type: ignore custom_toolsets: Optional[List[Path]] = opt_custom_toolsets, custom_runbooks: Optional[List[Path]] = opt_custom_runbooks, + allowed_builtin_toolsets: Optional[str] = opt_allowed_builtin_toolsets, max_steps: Optional[int] = opt_max_steps, verbose: Optional[List[bool]] = opt_verbose, # advanced options for this command @@ -893,6 +909,7 @@ def opsgenie( opsgenie_query=opsgenie_query, custom_toolsets_from_cli=custom_toolsets, custom_runbooks=custom_runbooks, + allowed_builtin_toolsets=allowed_builtin_toolsets, ) ai = config.create_console_issue_investigator() source = config.create_opsgenie_source() diff --git a/holmes/plugins/toolsets/__init__.py b/holmes/plugins/toolsets/__init__.py index 6932c82edf..6d161a5535 100644 --- a/holmes/plugins/toolsets/__init__.py +++ b/holmes/plugins/toolsets/__init__.py @@ -99,7 +99,18 @@ def load_python_toolsets(dal: Optional[SupabaseDal]) -> List[Toolset]: return toolsets -def load_builtin_toolsets(dal: Optional[SupabaseDal] = None) -> List[Toolset]: +def load_builtin_toolsets( + dal: Optional[SupabaseDal] = None, + allowed_builtin_toolsets: Optional[List[str]] = None, +) -> List[Toolset]: + """ + Load builtin toolsets from YAML files and Python modules. + + If allowed_builtin_toolsets is provided, filtering preserves the original + discovery order (loader order) rather than reordering by the allowlist. + This maintains consistent behavior where toolsets appear in the order + they were discovered during loading. + """ all_toolsets: List[Toolset] = [] logging.debug(f"loading toolsets from {THIS_DIR}") @@ -115,15 +126,36 @@ def load_builtin_toolsets(dal: Optional[SupabaseDal] = None) -> List[Toolset]: toolsets_from_file = load_toolsets_from_file(path, strict_check=True) all_toolsets.extend(toolsets_from_file) - all_toolsets.extend(load_python_toolsets(dal=dal)) # type: ignore + all_toolsets.extend(load_python_toolsets(dal=dal)) + + # Apply filter if specified + if allowed_builtin_toolsets is not None: + # Convert to set for O(1) lookup performance + requested_names = set(allowed_builtin_toolsets) + actual_names = {toolset.name for toolset in all_toolsets} + invalid_names = requested_names - actual_names + + if invalid_names: + logging.warning( + f"Unknown builtin toolsets specified: {sorted(invalid_names)}" + ) + + # Filter to only allowed toolsets using set for efficient lookup + # Note: This preserves the original discovery order rather than + # reordering by the allowlist + all_toolsets = [ + toolset + for toolset in all_toolsets + if toolset.name in requested_names # O(1) lookup + ] - # disable built-in toolsets by default, and the user can enable them explicitly in config. + # disable builtin toolsets by default, and the user can enable them explicitly in config. for toolset in all_toolsets: toolset.type = ToolsetType.BUILTIN - # dont' expose build-in toolsets path + # don't expose builtin toolsets path toolset.path = None - return all_toolsets # type: ignore + return all_toolsets def is_old_toolset_config( diff --git a/tests/core/test_config_allowed_builtin_toolsets.py b/tests/core/test_config_allowed_builtin_toolsets.py new file mode 100644 index 0000000000..8258d8bce0 --- /dev/null +++ b/tests/core/test_config_allowed_builtin_toolsets.py @@ -0,0 +1,131 @@ +from holmes.config import Config + + +def test_config_default_none(): + """Test that default value is None for backward compatibility""" + config = Config() + assert config.allowed_builtin_toolsets is None + + +def test_config_parse_comma_separated(tmp_path): + """Test parsing comma-separated string to list through actual config loading""" + # Create a minimal config file - can be empty since we're testing CLI option parsing + config_file = tmp_path / "config.yaml" + config_content = """ +# Empty config file - testing CLI argument parsing +model: "gpt-4o" +""" + config_file.write_text(config_content) + + # Pass the comma-separated string as a CLI kwarg to test the parsing logic + config = Config.load_from_file( + config_file, allowed_builtin_toolsets="kubernetes/core,prometheus/core" + ) + + # Assert that the parsing worked correctly + assert config.allowed_builtin_toolsets == ["kubernetes/core", "prometheus/core"] + + +def test_config_handle_whitespace(tmp_path): + """Test handling of whitespace and empty strings""" + # Create a minimal config file for testing CLI argument parsing + config_file = tmp_path / "config.yaml" + config_content = """ +# Test config for whitespace handling +model: "gpt-4o" +""" + config_file.write_text(config_content) + + test_cases = [ + ("kubernetes/core, prometheus/core", ["kubernetes/core", "prometheus/core"]), + ("kubernetes/core, ,prometheus/core", ["kubernetes/core", "prometheus/core"]), + (" kubernetes/core ", ["kubernetes/core"]), + ("", []), + ] + + for input_str, expected in test_cases: + # Use the public API to test the parsing logic + config = Config.load_from_file(config_file, allowed_builtin_toolsets=input_str) + assert config.allowed_builtin_toolsets == expected + + +def test_config_backward_compatibility(): + """Test that existing config loading works unchanged""" + config = Config() + assert hasattr(config, "allowed_builtin_toolsets") + assert config.allowed_builtin_toolsets is None + + +def test_config_load_from_file_with_allowed_toolsets(): + """Test that load_from_file correctly parses the allowed_builtin_toolsets option""" + # Test with comma-separated string + config = Config.load_from_file( + None, allowed_builtin_toolsets="kubernetes/core,prometheus/core" + ) + assert config.allowed_builtin_toolsets == ["kubernetes/core", "prometheus/core"] + + +def test_config_load_from_file_with_whitespace(): + """Test that load_from_file handles whitespace correctly""" + # Test with whitespace and empty strings + config = Config.load_from_file( + None, allowed_builtin_toolsets="kubernetes/core, prometheus/core, " + ) + assert config.allowed_builtin_toolsets == ["kubernetes/core", "prometheus/core"] + + +def test_config_load_from_file_with_empty_string(): + """Test that load_from_file handles empty string correctly""" + config = Config.load_from_file(None, allowed_builtin_toolsets="") + assert config.allowed_builtin_toolsets == [] + + +def test_config_load_from_file_without_option(): + """Test that load_from_file works when option is not provided""" + config = Config.load_from_file(None) + assert config.allowed_builtin_toolsets is None + + +def test_config_load_from_file_with_none(): + """Test that load_from_file handles None value correctly""" + config = Config.load_from_file(None, allowed_builtin_toolsets=None) + assert config.allowed_builtin_toolsets is None + + +def test_config_load_from_file_single_toolset(): + """Test parsing single toolset name""" + config = Config.load_from_file(None, allowed_builtin_toolsets="kubernetes/core") + assert config.allowed_builtin_toolsets == ["kubernetes/core"] + + +def test_config_load_from_file_complex_names(): + """Test parsing complex toolset names with special characters""" + config = Config.load_from_file( + None, allowed_builtin_toolsets="aws/ec2,grafana/loki,kubernetes/core" + ) + assert config.allowed_builtin_toolsets == [ + "aws/ec2", + "grafana/loki", + "kubernetes/core", + ] + + +def test_config_direct_assignment(): + """Test that the field can be assigned directly""" + config = Config(allowed_builtin_toolsets=["kubernetes/core", "prometheus/core"]) + assert config.allowed_builtin_toolsets == ["kubernetes/core", "prometheus/core"] + + +def test_config_field_type_checking(): + """Test that the field accepts correct types""" + # Test with None + config = Config(allowed_builtin_toolsets=None) + assert config.allowed_builtin_toolsets is None + + # Test with empty list + config = Config(allowed_builtin_toolsets=[]) + assert config.allowed_builtin_toolsets == [] + + # Test with list of strings + config = Config(allowed_builtin_toolsets=["test/toolset"]) + assert config.allowed_builtin_toolsets == ["test/toolset"] diff --git a/tests/core/test_toolset_manager_integration.py b/tests/core/test_toolset_manager_integration.py new file mode 100644 index 0000000000..3a4086f2c0 --- /dev/null +++ b/tests/core/test_toolset_manager_integration.py @@ -0,0 +1,250 @@ +import pytest +from holmes.core.toolset_manager import ToolsetManager +from holmes.config import Config +from holmes.core.tools import ToolsetType + + +def test_toolset_manager_filtering_integration(): + """Test end-to-end filtering through ToolsetManager""" + # Create config with filter + config = Config(allowed_builtin_toolsets=["kubernetes/core"]) + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + builtin_toolsets = [t for t in toolsets if t.type == ToolsetType.BUILTIN] + + # Should only contain allowed toolsets + assert len(builtin_toolsets) <= 1 # May be 0 if kubernetes/core not available + if builtin_toolsets: + assert builtin_toolsets[0].name == "kubernetes/core" + + +def test_custom_toolsets_unaffected(): + """Test that custom toolsets are not affected by filtering""" + config = Config(allowed_builtin_toolsets=["kubernetes/core"]) + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + custom_toolsets = [t for t in toolsets if t.type == ToolsetType.CUSTOMIZED] + + # Custom toolsets should be unaffected + # Compare with no-filter case + config_no_filter = Config() + manager_no_filter = ToolsetManager(config=config_no_filter) + toolsets_no_filter = manager_no_filter._list_all_toolsets() + custom_toolsets_no_filter = [ + t for t in toolsets_no_filter if t.type == ToolsetType.CUSTOMIZED + ] + + assert len(custom_toolsets) == len(custom_toolsets_no_filter) + + +def test_toolset_manager_no_filter(): + """Test ToolsetManager works normally when no filter specified""" + config = Config() # No filter + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + builtin_toolsets = [t for t in toolsets if t.type == ToolsetType.BUILTIN] + + # Should load all builtin toolsets + assert len(builtin_toolsets) > 0 # Assuming some builtin toolsets exist + + +def test_toolset_manager_empty_filter(): + """Test ToolsetManager with empty filter list""" + config = Config(allowed_builtin_toolsets=[]) + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + builtin_toolsets = [t for t in toolsets if t.type == ToolsetType.BUILTIN] + + # Should have no builtin toolsets + assert len(builtin_toolsets) == 0 + + +def test_toolset_manager_multiple_filters(): + """Test ToolsetManager with multiple allowed toolsets""" + # Get all toolsets first to select valid names + config_all = Config() + manager_all = ToolsetManager(config=config_all) + all_toolsets = manager_all._list_all_toolsets() + all_builtin = [t for t in all_toolsets if t.type == ToolsetType.BUILTIN] + + if len(all_builtin) < 2: + pytest.skip("Need at least 2 builtin toolsets for this test") + + # Select first two toolset names + target_names = [all_builtin[0].name, all_builtin[1].name] + + config = Config(allowed_builtin_toolsets=target_names) + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + builtin_toolsets = [t for t in toolsets if t.type == ToolsetType.BUILTIN] + + # Should have exactly 2 builtin toolsets + assert len(builtin_toolsets) == 2 + assert {t.name for t in builtin_toolsets} == set(target_names) + + +def test_toolset_manager_invalid_names(): + """Test ToolsetManager with invalid toolset names""" + config = Config(allowed_builtin_toolsets=["nonexistent/toolset"]) + manager = ToolsetManager(config=config) + + toolsets = manager._list_all_toolsets() + builtin_toolsets = [t for t in toolsets if t.type == ToolsetType.BUILTIN] + + # Should have no builtin toolsets (invalid names filtered out) + assert len(builtin_toolsets) == 0 + + +def test_toolset_manager_backward_compatibility(): + """Test that ToolsetManager maintains backward compatibility""" + # Test that ToolsetManager can be created without config parameter + manager_without_config = ToolsetManager() + toolsets_without_config = manager_without_config._list_all_toolsets() + + # Test that ToolsetManager with None config works the same + manager_with_none = ToolsetManager(config=None) + toolsets_with_none = manager_with_none._list_all_toolsets() + + # Should be identical + assert len(toolsets_without_config) == len(toolsets_with_none) + assert {t.name for t in toolsets_without_config} == { + t.name for t in toolsets_with_none + } + + +def test_config_toolset_manager_property(): + """Test that Config.toolset_manager property works with filtering""" + config = Config(allowed_builtin_toolsets=["kubernetes/core"]) + + # Access toolset_manager property + manager = config.toolset_manager + assert isinstance(manager, ToolsetManager) + assert manager._config is config + + # Test that _get_allowed_builtin_toolsets returns correct value + allowed = manager._get_allowed_builtin_toolsets() + assert allowed == ["kubernetes/core"] + + +def test_config_toolset_manager_no_filter(): + """Test that Config.toolset_manager works without filter""" + config = Config() + + # Access toolset_manager property + manager = config.toolset_manager + assert isinstance(manager, ToolsetManager) + assert manager._config is config + + # Test that _get_allowed_builtin_toolsets returns None + allowed = manager._get_allowed_builtin_toolsets() + assert allowed is None + + +def test_toolset_manager_caching(): + """Test that ToolsetManager is cached in Config""" + config = Config(allowed_builtin_toolsets=["kubernetes/core"]) + + # Access toolset_manager twice + manager1 = config.toolset_manager + manager2 = config.toolset_manager + + # Should be the same instance (cached) + assert manager1 is manager2 + + +def test_end_to_end_filtering_flow(): + """Test the complete end-to-end filtering flow""" + # This test verifies that the complete chain works: + # Config -> ToolsetManager -> load_builtin_toolsets -> filtered results + + # Get baseline - all toolsets + config_all = Config() + all_toolsets = config_all.toolset_manager._list_all_toolsets() + all_builtin_names = [t.name for t in all_toolsets if t.type == ToolsetType.BUILTIN] + + if not all_builtin_names: + pytest.skip("No builtin toolsets available") + + # Test filtering to single toolset + target_name = all_builtin_names[0] + config_filtered = Config(allowed_builtin_toolsets=[target_name]) + filtered_toolsets = config_filtered.toolset_manager._list_all_toolsets() + filtered_builtin = [t for t in filtered_toolsets if t.type == ToolsetType.BUILTIN] + + # Verify filtering worked + assert len(filtered_builtin) == 1 + assert filtered_builtin[0].name == target_name + + # Verify other toolset types are unaffected + all_non_builtin = [t for t in all_toolsets if t.type != ToolsetType.BUILTIN] + filtered_non_builtin = [ + t for t in filtered_toolsets if t.type != ToolsetType.BUILTIN + ] + + # Non-builtin toolsets should be unchanged + assert len(all_non_builtin) == len(filtered_non_builtin) + + +def test_console_toolsets_filtering(): + """Test that list_console_toolsets respects filtering""" + # Get all console toolsets without filtering + config_all = Config() + manager_all = ToolsetManager(config=config_all) + all_console_toolsets = manager_all.list_console_toolsets() + all_console_builtin = [ + t for t in all_console_toolsets if t.type == ToolsetType.BUILTIN + ] + + if not all_console_builtin: + pytest.skip("No builtin console toolsets available") + + # Filter to one toolset + target_name = all_console_builtin[0].name + config_filtered = Config(allowed_builtin_toolsets=[target_name]) + manager_filtered = ToolsetManager(config=config_filtered) + filtered_console_toolsets = manager_filtered.list_console_toolsets() + filtered_console_builtin = [ + t for t in filtered_console_toolsets if t.type == ToolsetType.BUILTIN + ] + + # Should only have the one allowed toolset + assert ( + len(filtered_console_builtin) <= 1 + ) # May be 0 if target doesn't have console tags + if filtered_console_builtin: + assert filtered_console_builtin[0].name == target_name + + +def test_server_toolsets_filtering(): + """Test that list_server_toolsets respects filtering""" + # Get all server toolsets without filtering + config_all = Config() + manager_all = ToolsetManager(config=config_all) + all_server_toolsets = manager_all.list_server_toolsets() + all_server_builtin = [ + t for t in all_server_toolsets if t.type == ToolsetType.BUILTIN + ] + + if not all_server_builtin: + pytest.skip("No builtin server toolsets available") + + # Filter to one toolset + target_name = all_server_builtin[0].name + config_filtered = Config(allowed_builtin_toolsets=[target_name]) + manager_filtered = ToolsetManager(config=config_filtered) + filtered_server_toolsets = manager_filtered.list_server_toolsets() + filtered_server_builtin = [ + t for t in filtered_server_toolsets if t.type == ToolsetType.BUILTIN + ] + + # Should only have the one allowed toolset + assert ( + len(filtered_server_builtin) <= 1 + ) # May be 0 if target doesn't have server tags + if filtered_server_builtin: + assert filtered_server_builtin[0].name == target_name diff --git a/tests/integration/test_allowed_builtin_toolsets_full.py b/tests/integration/test_allowed_builtin_toolsets_full.py new file mode 100644 index 0000000000..e534b7ab21 --- /dev/null +++ b/tests/integration/test_allowed_builtin_toolsets_full.py @@ -0,0 +1,256 @@ +import pytest +from unittest.mock import patch +from holmes.config import Config +from holmes.core.toolset_manager import ToolsetManager +from holmes.plugins.toolsets import load_builtin_toolsets +from holmes.core.tools import ToolsetType + + +class TestAllowedBuiltinToolsetsIntegration: + """Comprehensive integration tests for the complete feature""" + + def test_full_backward_compatibility(self): + """Test that all existing functionality works unchanged""" + # Test default behavior + config_default = Config() + manager_default = ToolsetManager(config=config_default) + toolsets_default = manager_default._list_all_toolsets() + + # Test explicit None + config_none = Config(allowed_builtin_toolsets=None) + manager_none = ToolsetManager(config=config_none) + toolsets_none = manager_none._list_all_toolsets() + + # Should be identical + assert len(toolsets_default) == len(toolsets_none) + assert {t.name for t in toolsets_default} == {t.name for t in toolsets_none} + + def test_error_handling(self): + """Test error handling for malformed input""" + # Test various edge cases that shouldn't crash + test_cases = [ + [], # Empty list + [""], # Empty string in list + [" "], # Whitespace only + ["invalid/toolset"], # Non-existent toolset + ["kubernetes/core", "", "prometheus/core"], # Mixed valid/invalid + ] + + for test_case in test_cases: + config = Config(allowed_builtin_toolsets=test_case) + manager = ToolsetManager(config=config) + # Should not raise exception + toolsets = manager._list_all_toolsets() + assert isinstance(toolsets, list) + + @pytest.mark.slow + def test_all_builtin_toolsets_individually(self): + """Test that each builtin toolset can be filtered individually""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("no builtin toolsets discovered") + + for toolset in all_toolsets: + config = Config(allowed_builtin_toolsets=[toolset.name]) + manager = ToolsetManager(config=config) + filtered_toolsets = manager._list_all_toolsets() + + builtin_filtered = [ + t for t in filtered_toolsets if t.type == ToolsetType.BUILTIN + ] + assert len(builtin_filtered) == 1 + assert builtin_filtered[0].name == toolset.name + + def test_config_validation(self): + """Test that Config field validation works correctly""" + # Test that empty strings and whitespace are filtered out + config = Config( + allowed_builtin_toolsets=["kubernetes/core", "", " ", "prometheus/core"] + ) + # Validation should filter out empty/whitespace strings + expected = ["kubernetes/core", "prometheus/core"] + assert config.allowed_builtin_toolsets == expected + + # Test that all empty results in empty list (not None) + config_empty = Config(allowed_builtin_toolsets=["", " "]) + assert config_empty.allowed_builtin_toolsets == [] + + def test_warning_for_invalid_toolsets(self): + """Test that warnings are logged for invalid toolset names""" + # Test with invalid toolset names + with patch("holmes.plugins.toolsets.logging.warning") as mock_warning: + config = Config( + allowed_builtin_toolsets=["nonexistent/toolset", "another/fake"] + ) + manager = ToolsetManager(config=config) + manager._list_all_toolsets() # This will trigger the filtering + + # Should have called warning about invalid names + mock_warning.assert_called() + # Check the last call for the expected message content + last_call_args = mock_warning.call_args[0][0] + assert "Unknown builtin toolsets specified" in last_call_args + assert "nonexistent/toolset" in last_call_args + assert "another/fake" in last_call_args + + def test_warning_partial_invalid_toolsets(self): + """Test warnings when some toolsets are valid and some invalid""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + valid_name = all_toolsets[0].name + invalid_names = ["nonexistent/toolset"] + + with patch("holmes.plugins.toolsets.logging.warning") as mock_warning: + config = Config(allowed_builtin_toolsets=[valid_name] + invalid_names) + manager = ToolsetManager(config=config) + filtered_toolsets = manager._list_all_toolsets() + + # Should warn about invalid names + assert mock_warning.called + warning_message = mock_warning.call_args[0][0] + assert "Unknown builtin toolsets specified" in warning_message + assert "nonexistent/toolset" in warning_message + + # But should still include valid toolset + builtin_filtered = [ + t for t in filtered_toolsets if t.type == ToolsetType.BUILTIN + ] + assert len(builtin_filtered) == 1 + assert builtin_filtered[0].name == valid_name + + def test_no_warning_for_valid_toolsets(self): + """Test that no warnings are logged for valid toolset names""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + valid_name = all_toolsets[0].name + + with patch("logging.warning") as mock_warning: + config = Config(allowed_builtin_toolsets=[valid_name]) + manager = ToolsetManager(config=config) + manager._list_all_toolsets() + + # Should not have called warning about unknown toolsets specifically + assert not any( + "Unknown builtin toolsets specified" in str(args[0]) + if args + else False + or "Unknown builtin toolsets specified" in kwargs.get("msg", "") + for args, kwargs in mock_warning.call_args_list + ) + + def test_filtering_preserves_toolset_order(self): + """Test that filtering maintains the original order of toolsets""" + all_toolsets = load_builtin_toolsets() + if len(all_toolsets) < 3: + pytest.skip("Need at least 3 toolsets for order test") + + # Select first and third toolsets + target_names = [all_toolsets[0].name, all_toolsets[2].name] + + config = Config(allowed_builtin_toolsets=target_names) + manager = ToolsetManager(config=config) + filtered_toolsets = manager._list_all_toolsets() + builtin_filtered = [ + t for t in filtered_toolsets if t.type == ToolsetType.BUILTIN + ] + + # Should maintain original order + assert len(builtin_filtered) == 2 + assert builtin_filtered[0].name == all_toolsets[0].name + assert builtin_filtered[1].name == all_toolsets[2].name + + def test_config_cli_parsing_integration(self): + """Test integration between CLI parsing and Config validation""" + # Test that CLI parsing + validation works together + config = Config.load_from_file( + None, allowed_builtin_toolsets="kubernetes/core, , prometheus/core, " + ) + + # Should parse comma-separated string and validate + expected = ["kubernetes/core", "prometheus/core"] + assert config.allowed_builtin_toolsets == expected + + def test_empty_filter_behavior(self): + """Test behavior when filter results in empty list""" + config = Config(allowed_builtin_toolsets=[]) + manager = ToolsetManager(config=config) + filtered_toolsets = manager._list_all_toolsets() + builtin_filtered = [ + t for t in filtered_toolsets if t.type == ToolsetType.BUILTIN + ] + + # Should have no builtin toolsets + assert len(builtin_filtered) == 0 + + # But other types should still be present + non_builtin = [t for t in filtered_toolsets if t.type != ToolsetType.BUILTIN] + # The count depends on what's configured, but should not crash + assert isinstance(non_builtin, list) + + def test_toolset_manager_caching_with_filter(self): + """Test that Config caches ToolsetManager correctly with filtering""" + config = Config(allowed_builtin_toolsets=["kubernetes/core"]) + + # Access toolset_manager multiple times + manager1 = config.toolset_manager + manager2 = config.toolset_manager + + # Should be same instance (cached) + assert manager1 is manager2 + + # Both should have same config reference + assert manager1._config is config + assert manager2._config is config + + def test_console_and_server_toolset_filtering(self): + """Test that console and server toolset methods respect filtering""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + target_name = all_toolsets[0].name + config = Config(allowed_builtin_toolsets=[target_name]) + manager = ToolsetManager(config=config) + + # Test console toolsets + console_toolsets = manager.list_console_toolsets() + console_builtin = [t for t in console_toolsets if t.type == ToolsetType.BUILTIN] + + # Should only include allowed toolsets (filtered by both tags and allowed list) + for toolset in console_builtin: + assert toolset.name == target_name + + # Test server toolsets + server_toolsets = manager.list_server_toolsets() + server_builtin = [t for t in server_toolsets if t.type == ToolsetType.BUILTIN] + + # Should only include allowed toolsets (filtered by both tags and allowed list) + for toolset in server_builtin: + assert toolset.name == target_name + + def test_feature_robustness(self): + """Test that feature is robust to various edge cases""" + edge_cases = [ + None, # None value + [], # Empty list + [""], # List with empty string + ["kubernetes/core"] * 5, # Duplicate names + ["KUBERNETES/CORE"], # Wrong case (should not match) + ["kubernetes/core/extra"], # Extra path components + ] + + for case in edge_cases: + try: + config = Config(allowed_builtin_toolsets=case) + manager = ToolsetManager(config=config) + toolsets = manager._list_all_toolsets() + # Should not crash and return a list + assert isinstance(toolsets, list) + except Exception as e: + pytest.fail( + f"Feature should be robust to edge case {case}, but got: {e}" + ) diff --git a/tests/plugins/test_toolsets_filtering.py b/tests/plugins/test_toolsets_filtering.py new file mode 100644 index 0000000000..025ca02c36 --- /dev/null +++ b/tests/plugins/test_toolsets_filtering.py @@ -0,0 +1,178 @@ +import pytest +from holmes.plugins.toolsets import load_builtin_toolsets + + +def test_load_builtin_toolsets_no_filter(): + """Test that None filter preserves existing behavior""" + toolsets_none = load_builtin_toolsets(allowed_builtin_toolsets=None) + toolsets_default = load_builtin_toolsets() # Default behavior + + assert len(toolsets_none) == len(toolsets_default) + assert {t.name for t in toolsets_none} == {t.name for t in toolsets_default} + + +def test_load_builtin_toolsets_with_single_filter(): + """Test filtering to single toolset""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + first_toolset_name = all_toolsets[0].name + filtered = load_builtin_toolsets(allowed_builtin_toolsets=[first_toolset_name]) + + assert len(filtered) == 1 + assert filtered[0].name == first_toolset_name + + +def test_load_builtin_toolsets_with_multiple_filters(): + """Test filtering to multiple toolsets""" + all_toolsets = load_builtin_toolsets() + if len(all_toolsets) < 2: + pytest.skip("Need at least 2 builtin toolsets") + + # Use reverse order to test that discovery order is preserved + target_names = [all_toolsets[1].name, all_toolsets[0].name] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=target_names) + + assert len(filtered) == 2 + # Should preserve discovery order, not filter order + assert [t.name for t in filtered] == [all_toolsets[0].name, all_toolsets[1].name] + + +def test_load_builtin_toolsets_empty_filter(): + """Test that empty list results in no toolsets""" + filtered = load_builtin_toolsets(allowed_builtin_toolsets=[]) + assert len(filtered) == 0 + + +def test_load_builtin_toolsets_invalid_names(): + """Test that invalid toolset names are handled gracefully""" + filtered = load_builtin_toolsets(allowed_builtin_toolsets=["nonexistent/toolset"]) + assert len(filtered) == 0 + + +def test_load_builtin_toolsets_mixed_valid_invalid(): + """Test mix of valid and invalid names""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + valid_name = all_toolsets[0].name + mixed_names = [valid_name, "nonexistent/toolset"] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=mixed_names) + + assert len(filtered) == 1 + assert filtered[0].name == valid_name + + +def test_load_builtin_toolsets_preserves_toolset_properties(): + """Test that filtered toolsets maintain their properties""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + target_toolset = all_toolsets[0] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=[target_toolset.name]) + + assert len(filtered) == 1 + filtered_toolset = filtered[0] + + # Check that properties are preserved + assert filtered_toolset.name == target_toolset.name + assert filtered_toolset.type == target_toolset.type + assert filtered_toolset.path == target_toolset.path + + +def test_load_builtin_toolsets_with_dal_parameter(): + """Test that dal parameter works with filtering""" + # Test with dal=None and filtering + filtered = load_builtin_toolsets(dal=None, allowed_builtin_toolsets=[]) + assert len(filtered) == 0 + + # Test that dal parameter doesn't interfere with filtering + all_toolsets = load_builtin_toolsets(dal=None) + if all_toolsets: + target_name = all_toolsets[0].name + filtered_with_dal = load_builtin_toolsets( + dal=None, allowed_builtin_toolsets=[target_name] + ) + assert len(filtered_with_dal) == 1 + assert filtered_with_dal[0].name == target_name + + +def test_load_builtin_toolsets_case_sensitive_matching(): + """Test that toolset name matching is case-sensitive""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + original_name = all_toolsets[0].name + uppercase_name = original_name.upper() + + # Should not match if case is different + if uppercase_name != original_name: # Only test if case actually differs + filtered = load_builtin_toolsets(allowed_builtin_toolsets=[uppercase_name]) + assert len(filtered) == 0 + + +def test_load_builtin_toolsets_exact_name_matching(): + """Test that toolset name matching is exact (not substring)""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + original_name = all_toolsets[0].name + partial_name = ( + original_name[: len(original_name) // 2] + if len(original_name) > 1 + else original_name + ) + + # Should not match partial name if it's different from full name + if partial_name != original_name: + filtered = load_builtin_toolsets(allowed_builtin_toolsets=[partial_name]) + assert len(filtered) == 0 + + +def test_load_builtin_toolsets_all_valid_names(): + """Test filtering with all valid toolset names (should return all toolsets)""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + all_names = [toolset.name for toolset in all_toolsets] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=all_names) + + assert len(filtered) == len(all_toolsets) + assert {t.name for t in filtered} == {t.name for t in all_toolsets} + + +def test_load_builtin_toolsets_duplicate_names_in_filter(): + """Test that duplicate names in filter list don't cause issues""" + all_toolsets = load_builtin_toolsets() + if not all_toolsets: + pytest.skip("No builtin toolsets available") + + target_name = all_toolsets[0].name + duplicate_filter = [target_name, target_name, target_name] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=duplicate_filter) + + # Should still return only one instance of the toolset + assert len(filtered) == 1 + assert filtered[0].name == target_name + + +def test_load_builtin_toolsets_maintains_order(): + """Test that filtering maintains the original order of toolsets""" + all_toolsets = load_builtin_toolsets() + if len(all_toolsets) < 3: + pytest.skip("Need at least 3 builtin toolsets") + + # Select third then first toolsets (reverse of discovery order) + target_names = [all_toolsets[2].name, all_toolsets[0].name] + filtered = load_builtin_toolsets(allowed_builtin_toolsets=target_names) + + assert len(filtered) == 2 + # Should preserve discovery order, not allowlist order + assert filtered[0].name == all_toolsets[0].name + assert filtered[1].name == all_toolsets[2].name