diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index 655ac9a..78cd750 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -22,12 +22,6 @@ ) from galileo_core.schemas.logging.step import StepType from galileo_core.schemas.logging.trace import Trace -from galileo_core.schemas.protect.execution_status import ExecutionStatus -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.request import Request -from galileo_core.schemas.protect.response import Response -from galileo_core.schemas.protect.ruleset import Ruleset -from galileo_core.schemas.protect.stage import StageType from splunk_ao.agent_control import AgentControlTarget, AgentControlTargetUnresolvedError, get_agent_control_target from splunk_ao.collaborator import Collaborator, CollaboratorRole from splunk_ao.configuration import Configuration @@ -54,7 +48,6 @@ from splunk_ao.model import Model from splunk_ao.project import Project from splunk_ao.prompt import Prompt -from splunk_ao.protect import ainvoke_protect, invoke_protect from splunk_ao.provider import AnthropicProvider, AzureProvider, BedrockProvider, OpenAIProvider, Provider from splunk_ao.schema.message import Message from splunk_ao.schema.metrics import SplunkAOMetrics @@ -67,13 +60,6 @@ SplunkAOFutureError, ValidationError, ) -from splunk_ao.stages import ( - create_protect_stage, - get_protect_stage, - pause_protect_stage, - resume_protect_stage, - update_protect_stage, -) from splunk_ao.tracing import get_tracing_headers from splunk_ao.types import MetricSpec from splunk_ao.utils.log_config import enable_console_logging @@ -104,7 +90,6 @@ "ControlSpan", "Dataset", "Document", - "ExecutionStatus", "Experiment", "ForbiddenError", "Integration", @@ -119,17 +104,13 @@ "Model", "NotFoundError", "OpenAIProvider", - "Payload", "Project", "Prompt", "Provider", "RateLimitError", - "Request", "ResourceConflictError", "ResourceNotFoundError", - "Response", "RetrieverSpan", - "Ruleset", "ServerError", "Session", "Span", @@ -141,7 +122,6 @@ "SplunkAOLoggerException", "SplunkAOMetric", "SplunkAOMetrics", - "StageType", "StepType", "StepWithChildSpans", "SyncState", @@ -151,22 +131,15 @@ "Trace", "ValidationError", "WorkflowSpan", - "ainvoke_protect", "create_api_key", - "create_protect_stage", "delete_api_key", "enable_console_logging", "splunk_ao_context", "get_agent_control_target", - "get_protect_stage", "get_tracing_headers", - "invoke_protect", "is_dependency_available", "list_api_keys", "log", - "pause_protect_stage", - "resume_protect_stage", "setup_agent_control_bridge", "start_session", - "update_protect_stage", ] diff --git a/src/splunk_ao/constants/protect.py b/src/splunk_ao/constants/protect.py deleted file mode 100644 index 629071b..0000000 --- a/src/splunk_ao/constants/protect.py +++ /dev/null @@ -1 +0,0 @@ -TIMEOUT_SECS = 10 diff --git a/src/splunk_ao/handlers/langchain/tool.py b/src/splunk_ao/handlers/langchain/tool.py deleted file mode 100644 index 3ac89de..0000000 --- a/src/splunk_ao/handlers/langchain/tool.py +++ /dev/null @@ -1,164 +0,0 @@ -from collections.abc import Sequence - -from langchain_core.runnables.base import Runnable -from langchain_core.tools import BaseTool -from pydantic import UUID4, BaseModel, ConfigDict, Field - -from galileo_core.schemas.protect.execution_status import ExecutionStatus -from galileo_core.schemas.protect.payload import Payload as CorePayload -from galileo_core.schemas.protect.response import Response -from galileo_core.schemas.protect.ruleset import Ruleset -from splunk_ao.constants.protect import TIMEOUT_SECS -from splunk_ao.protect import ainvoke_protect, invoke_protect -from splunk_ao.utils.log_config import get_logger - -logger = get_logger(__name__) - - -class ProtectToolInputSchema(BaseModel): - """Input schema for the ProtectTool.""" - - input: str | None = None - output: str | None = None - - -class ProtectTool(BaseTool): - """A LangChain tool that wraps the Galileo Protect API. - - This tool allows you to integrate Galileo Protect into your LangChain applications - to scan for and mitigate harmful content in both inputs and outputs. It can be - configured with specific rulesets and linked to a Galileo project and stage for - monitoring and management. - - It accepts an input and/or output as arguments (parsed by a Pydantic model) and returns a JSON-serialized `Response` object from the Protect API. - - Attributes - ---------- - prioritized_rulesets: - An optional sequence of `Ruleset` objects to apply. - project_id: - The UUID of the Galileo project this tool is associated with. - project_name: - The name of the Galileo project. - stage_name: - The name of the Protect stage to use for this tool. - stage_id: - The UUID of the Protect stage. - stage_version: - The version of the Protect stage to use. - timeout: - The timeout in seconds for the API request. - """ - - name: str = "GalileoProtect" - description: str = ( - "Protect your LLM applications from harmful content using Galileo Protect. " - "This tool is a wrapper around Galileo's Protect API, can be used to scan text " - "for harmful content, and can be used to trigger actions based on the results." - "The tool can be used on the input text or output text, and can be configured " - "with a set of rulesets to evaluate on." - ) - args_schema: type[BaseModel] = ProtectToolInputSchema - - prioritized_rulesets: Sequence[Ruleset] | None = None - project_id: UUID4 | None = None - project_name: str | None = None - stage_name: str | None = None - stage_id: UUID4 | None = None - stage_version: int | None = None - timeout: float = TIMEOUT_SECS - - def _run(self, input: str | None = None, output: str | None = None) -> str: - """ - Apply the tool synchronously. - - We serialize the response to JSON because that's what `langchain_core` expects - for tools. - """ - api_payload = CorePayload(input=input, output=output) - api_response = invoke_protect( - payload=api_payload, - prioritized_rulesets=self.prioritized_rulesets, - project_name=self.project_name, - project_id=self.project_id, - stage_name=self.stage_name, - stage_id=self.stage_id, - stage_version=self.stage_version, - timeout=self.timeout, - ) - return api_response.model_dump_json() - - async def _arun(self, input: str | None = None, output: str | None = None) -> str: - """ - Apply the tool asynchronously. - - We serialize the response to JSON because that's what `langchain_core` expects - for tools. - """ - api_payload = CorePayload(input=input, output=output) - api_response = await ainvoke_protect( - payload=api_payload, - prioritized_rulesets=self.prioritized_rulesets, - project_name=self.project_name, - project_id=self.project_id, - stage_name=self.stage_name, - stage_id=self.stage_id, - stage_version=self.stage_version, - timeout=self.timeout, - ) - return api_response.model_dump_json() - - -class ProtectParser(BaseModel): - """A LangChain runnable that parses and routes the output of a ``ProtectTool``. - - If the Protect API response is 'triggered', it returns the response text. - Otherwise, it invokes a fallback chain. - - Attributes - ---------- - chain: The ``Runnable`` to invoke if the Protect invocation is not triggered. - ignore_trigger: If True, always invoke the fallback chain. - echo_output: If True, print the raw Protect API response to the console. - """ - - chain: Runnable = Field(..., description="The chain to trigger if the Protect invocation is not triggered.") - ignore_trigger: bool = Field( - default=False, - description="Ignore the status of the Protect invocation and always trigger the rest of the chain.", - ) - echo_output: bool = Field(default=False, description="Echo the output of the Protect invocation.") - - model_config = ConfigDict(arbitrary_types_allowed=True) - - def parser(self, response_raw_json: str) -> str: - """Parses the JSON response from ``ProtectTool`` and decides the execution path. - - If the response status is 'triggered', the response text is returned. Otherwise, - the fallback chain is invoked with the text. - - If JSON parsing fails, it assumes the input is not from ``ProtectTool`` and - invokes the fallback chain directly with the raw input. - - Parameters - ---------- - response_raw_json: str - Expects the output from the ``ProtectTool``. - - Returns - ------- - str - The text from the Protect response if triggered, or the result of invoking - the fallback chain. - """ - try: - response = Response.model_validate_json(response_raw_json) - except Exception: - return self.chain.invoke(response_raw_json) - text = response.text - - if self.echo_output: - logger.debug(f"> Raw response: {text}") - if response.status == ExecutionStatus.triggered and not self.ignore_trigger: - return text - return self.chain.invoke(text) diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index c2494cb..5879c93 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -33,8 +33,6 @@ ) from galileo_core.schemas.logging.step import BaseStep, Metrics, StepType from galileo_core.schemas.logging.trace import Trace -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.response import Response from galileo_core.schemas.shared.traces_logger import TracesLogger from splunk_ao.config import SplunkAOConfig from splunk_ao.constants import LoggerModeType @@ -1578,85 +1576,6 @@ def add_tool_span( return span - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def add_protect_span( - self, - payload: Payload, - redacted_payload: Payload | None = None, - response: Response | None = None, - redacted_response: Response | None = None, - created_at: datetime | None = None, - metadata: dict[str, MetadataValue] | None = None, - tags: list[str] | None = None, - status_code: int | None = None, - step_number: int | None = None, - ) -> ToolSpan: - """ - Add a new Protect tool span to the current parent. - - Parameters - ---------- - payload: Payload - Input to the node. This is the input to the Protect `invoke` method. - Expected format: Payload object with input_ and/or output attributes. - Example: `Payload(input_="User input text", output="Model output text")` - redacted_payload: Optional[Payload] - Input that removes any sensitive information (redacted input to the node). - Same format as payload parameter. - response: Optional[Response] - Output of the node. This is the output from the Protect `invoke` method. - Expected format: Response object with text, trace_metadata, and status. - Example: `Response(text="Processed text", status=ExecutionStatus.triggered)` - redacted_response: Optional[Response] - Output that removes any sensitive information (redacted output of the node). - Same format as response parameter. - created_at: Optional[datetime] - Timestamp of the span's creation. - metadata: Optional[dict[str, str]] - Metadata associated with this span. - Expected format: `{"key1": "value1", "key2": "value2"}` - tags: Optional[list[str]] - Tags associated with this span. - Expected format: `["tag1", "tag2", "tag3"]` - status_code: Optional[int] - Status code of the node execution. - Expected values: 200 (success), 400 (client error), 500 (server error) - step_number: Optional[int] - Step number of the span. - - Returns - ------- - ToolSpan - The created Protect tool span. - """ - # Auto-convert non-string metadata values to strings - if metadata: - metadata = {k: SplunkAOLogger._convert_metadata_value(v) for k, v in metadata.items()} - - kwargs = { - "input": json.dumps(payload.model_dump(mode="json")), - "redacted_input": json.dumps(redacted_payload.model_dump(mode="json")) if redacted_payload else None, - "output": json.dumps(response.model_dump(mode="json")) if response else None, - "redacted_output": json.dumps(redacted_response.model_dump(mode="json")) if redacted_response else None, - "name": "GalileoProtect", - "duration_ns": response.trace_metadata.response_at - response.trace_metadata.received_at - if response - else None, - "created_at": created_at, - "user_metadata": metadata, - "tags": tags, - "status_code": status_code, - "step_number": step_number, - "id": uuid.uuid4(), - } - span = super().add_tool_span(**kwargs) - - if self.mode == "distributed": - self._ingest_step_streaming(span) - - return span - def _attach_parentable_span(self, span: StepWithChildSpans, status_code: int | None = None) -> StepWithChildSpans: parent = self.current_parent() span._parent = parent diff --git a/src/splunk_ao/protect.py b/src/splunk_ao/protect.py deleted file mode 100644 index ed16d0c..0000000 --- a/src/splunk_ao/protect.py +++ /dev/null @@ -1,187 +0,0 @@ -from collections.abc import Sequence - -from pydantic import UUID4 - -from galileo.resources.api.protect import invoke_protect_invoke_post -from galileo.resources.models.http_validation_error import HTTPValidationError -from galileo.resources.models.protect_request import ProtectRequest as APIRequest -from galileo.resources.models.protect_response import ProtectResponse as APIResponse -from galileo_core.helpers.execution import async_run -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.request import Request -from galileo_core.schemas.protect.response import Response -from galileo_core.schemas.protect.ruleset import Ruleset -from splunk_ao.config import SplunkAOConfig -from splunk_ao.constants.protect import TIMEOUT_SECS - - -class Protect: - config: SplunkAOConfig - - def __init__(self) -> None: - self.config = SplunkAOConfig.get() - - async def ainvoke( - self, - payload: Payload, - prioritized_rulesets: Sequence[Ruleset] | None = None, - project_id: UUID4 | None = None, - project_name: str | None = None, - stage_id: UUID4 | None = None, - stage_name: str | None = None, - stage_version: int | None = None, - timeout: float = TIMEOUT_SECS, - metadata: dict[str, str] | None = None, - headers: dict[str, str] | None = None, - ) -> Response | HTTPValidationError | None: - request = Request( - payload=payload, - prioritized_rulesets=prioritized_rulesets or [], - project_id=str(project_id) if project_id is not None else None, - project_name=project_name, - stage_id=str(stage_id) if stage_id is not None else None, - stage_name=stage_name, - stage_version=stage_version, - timeout=timeout, - metadata=metadata, - headers=headers, - ) - request_dict = request.model_dump(mode="json") - request_dict["prioritized_rulesets"] = request_dict.pop("rulesets", []) - body = APIRequest.from_dict(request_dict) - - response: APIResponse | HTTPValidationError | None = await invoke_protect_invoke_post.asyncio( - client=self.config.api_client, body=body - ) - - if isinstance(response, APIResponse): - return Response.model_validate(response.to_dict()) - return response - - -async def ainvoke_protect( - payload: Payload, - prioritized_rulesets: Sequence[Ruleset] | None = None, - project_id: UUID4 | None = None, - project_name: str | None = None, - stage_id: UUID4 | None = None, - stage_name: str | None = None, - stage_version: int | None = None, - timeout: float = TIMEOUT_SECS, - metadata: dict[str, str] | None = None, - headers: dict[str, str] | None = None, -) -> Response | HTTPValidationError | None: - """Asynchronously invoke Protect with the given payload. - - If using the local stage, the prioritized rulesets should be provided to ensure the - correct rulesets are used for processing. If using a central stage, the rulesets - will be fetched from the existing stage definition. - - Project ID and stage name, or stage ID should be provided for all invocations. - - Parameters - ---------- - payload - Payload to be processed. - prioritized_rulesets - Prioritized rulesets to be used for processing. - These should only be provided if using a local stage. Defaults to an - empty list if None. - project_id - ID of the project. - project_name - Name of the project. - stage_id - ID of the stage. - stage_name - Name of the stage. - stage_version - Version of the stage. - timeout - Timeout for the request in seconds. Defaults to TIMEOUT_SECS. - metadata - Metadata to be added when responding. - headers - Headers to be added to the response. - - Returns - ------- - Protect invoke results. - """ - return await Protect().ainvoke( - payload=payload, - prioritized_rulesets=prioritized_rulesets, - project_id=project_id, - project_name=project_name, - stage_id=stage_id, - stage_name=stage_name, - stage_version=stage_version, - timeout=timeout, - metadata=metadata, - headers=headers, - ) - - -def invoke_protect( - payload: Payload, - prioritized_rulesets: Sequence[Ruleset] | None = None, - project_id: UUID4 | None = None, - project_name: str | None = None, - stage_id: UUID4 | None = None, - stage_name: str | None = None, - stage_version: int | None = None, - timeout: float = TIMEOUT_SECS, - metadata: dict[str, str] | None = None, - headers: dict[str, str] | None = None, -) -> Response | HTTPValidationError | None: - """Invoke Protect with the given payload. - - If using the local stage, the prioritized rulesets should be provided to ensure the - correct rulesets are used for processing. If using a central stage, the rulesets - will be fetched from the existing stage definition. - - Project ID and stage name, or stage ID should be provided for all invocations. - - Parameters - ---------- - payload - Payload to be processed. - prioritized_rulesets - Prioritized rulesets to be used for processing. - These should only be provided if using a local stage. Defaults to an - empty list if None. - project_id - ID of the project. - project_name - Name of the project. - stage_id - ID of the stage. - stage_name - Name of the stage. - stage_version - Version of the stage. - timeout - Timeout for the request in seconds. Defaults to TIMEOUT_SECS. - metadata - Metadata to be added when responding. - headers - Headers to be added to the response. - - Returns - ------- - Protect invoke results. - """ - return async_run( - ainvoke_protect( - payload=payload, - prioritized_rulesets=prioritized_rulesets, - project_id=project_id, - project_name=project_name, - stage_id=stage_id, - stage_name=stage_name, - stage_version=stage_version, - timeout=timeout, - metadata=metadata, - headers=headers, - ) - ) diff --git a/src/splunk_ao/stages.py b/src/splunk_ao/stages.py deleted file mode 100644 index 9ccbfb5..0000000 --- a/src/splunk_ao/stages.py +++ /dev/null @@ -1,346 +0,0 @@ -from collections.abc import Sequence - -from pydantic import UUID4 - -from galileo.resources.api.protect import ( - create_stage_projects_project_id_stages_post, - get_stage_projects_project_id_stages_get, - pause_stage_projects_project_id_stages_stage_id_put, - update_stage_projects_project_id_stages_stage_id_post, -) -from galileo.resources.models.rulesets_mixin import RulesetsMixin as APIRulesetsMixin -from galileo.resources.models.stage_db import StageDB as APIStageDB -from galileo.resources.models.stage_with_rulesets import StageWithRulesets as APIStageWithRulesets -from galileo.resources.types import UNSET -from galileo_core.schemas.protect.ruleset import Ruleset, RulesetsMixin -from galileo_core.schemas.protect.stage import StageDB, StageType, StageWithRulesets -from galileo_core.utils.name import ts_name -from splunk_ao.config import SplunkAOConfig -from splunk_ao.projects import Projects - - -def _get_validated_project_id(project_id: str | UUID4 | None = None, project_name: str | None = None) -> str: - """Resolves project ID from either project_id or project_name.""" - # If the project Id is a UUID4, convert to string. - if project_id is not None and type(project_id).__name__ == "UUID": - project_id = str(project_id) - - project = Projects().get_with_env_fallbacks(name=project_name, id=project_id) - if not project: - raise ValueError(f"Project with name '{project_name}' not found.") - return str(project.id) - - -def _get_stage_id( - stage_id: str | UUID4 | None = None, stage_name: str | None = None, project_id: str | UUID4 | None = None -) -> str: - """ - Resolves stage ID from either stage_id or stage_name. - If stage_name is provided, it will look up the stage within the specified project. - """ - if stage_id: - return str(stage_id) - if stage_name: - stage = Stages().get(project_id=project_id, stage_name=stage_name) - if not stage: - raise ValueError(f"Stage with name '{stage_name}' not found.") - return str(stage.id) - raise ValueError("Either stage_id or stage_name must be provided.") - - -class Stages: - config: SplunkAOConfig - - def __init__(self) -> None: - self.config = SplunkAOConfig.get() - - def create( - self, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - name: str | None = None, - stage_type: StageType = StageType.local, - pause: bool = False, - prioritized_rulesets: Sequence[Ruleset] | None = None, - description: str | None = None, - ) -> StageDB: - actual_project_id: str = _get_validated_project_id(project_id=project_id, project_name=project_name) - - actual_name = name or ts_name("stage") - - request = StageWithRulesets( - name=actual_name, - project_id=str(actual_project_id), - type=stage_type, - paused=pause, - description=description, - prioritized_rulesets=prioritized_rulesets or [], - ) - - request_dict = request.model_dump(mode="json") - request_dict["prioritized_rulesets"] = request_dict.pop("rulesets", []) - body = APIStageWithRulesets.from_dict(request_dict) - - response = create_stage_projects_project_id_stages_post.sync( - project_id=str(actual_project_id), client=self.config.api_client, body=body - ) - - if isinstance(response, APIStageDB): - return StageDB.model_validate(response.to_dict()) - return response - - def get( - self, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - ) -> StageDB: - actual_project_id: str = _get_validated_project_id(project_id=project_id, project_name=project_name) - - if not stage_id and not stage_name: - raise ValueError("Either stage_id or stage_name must be provided.") - - response = get_stage_projects_project_id_stages_get.sync( - project_id=actual_project_id, - stage_id=str(stage_id) if stage_id else UNSET, - stage_name=stage_name if stage_name else UNSET, - client=self.config.api_client, - ) - if isinstance(response, APIStageDB): - return StageDB.model_validate(response.to_dict()) - return response - - def update( - self, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - prioritized_rulesets: Sequence[Ruleset] | None = None, - ) -> StageDB: - actual_project_id: str = _get_validated_project_id(project_id=project_id, project_name=project_name) - - actual_stage_id: str = _get_stage_id(stage_id=stage_id, stage_name=stage_name, project_id=actual_project_id) - - request = RulesetsMixin(prioritized_rulesets=prioritized_rulesets or []) - request_dict = request.model_dump() - request_dict["prioritized_rulesets"] = request_dict.pop("rulesets", []) - body = APIRulesetsMixin.from_dict(request_dict) - - response = update_stage_projects_project_id_stages_stage_id_post.sync( - project_id=actual_project_id, stage_id=actual_stage_id, client=self.config.api_client, body=body - ) - if isinstance(response, APIStageDB): - return StageDB.model_validate(response.to_dict()) - return response - - def _set_pause_state( - self, - pause_flag: bool, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - ) -> StageDB: - """Sets the pause state of a stage.""" - actual_project_id: str = _get_validated_project_id(project_id=project_id, project_name=project_name) - - actual_stage_id: str = _get_stage_id(stage_id=stage_id, stage_name=stage_name, project_id=actual_project_id) - - response = pause_stage_projects_project_id_stages_stage_id_put.sync( - project_id=actual_project_id, stage_id=actual_stage_id, client=self.config.api_client, pause=pause_flag - ) - if isinstance(response, APIStageDB): - return StageDB.model_validate(response.to_dict()) - return response - - def pause( - self, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - ) -> StageDB: - return self._set_pause_state( - project_id=project_id, project_name=project_name, stage_id=stage_id, stage_name=stage_name, pause_flag=True - ) - - def resume( - self, - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - ) -> StageDB: - return self._set_pause_state( - project_id=project_id, project_name=project_name, stage_id=stage_id, stage_name=stage_name, pause_flag=False - ) - - -def create_protect_stage( - project_id: str | UUID4 | None = None, - project_name: str | None = None, - name: str | None = None, - stage_type: StageType = StageType.local, - pause: bool = False, - prioritized_rulesets: Sequence[Ruleset] | None = None, - description: str | None = None, -) -> StageDB | None: - """Creates a new stage. - - Parameters - ---------- - project_id - The ID of the project. - project_name - Name of the project. If `project_id` is not provided, - this will be used to look up the project. - name - Name of the stage. Defaults to a generated name. - stage_type - Type of the stage. - pause - Whether the stage should be created in a paused state. - prioritized_rulesets - List of rulesets for the stage. - description - Description for the stage. - - Returns - ------- - The newly created Stage. - """ - return Stages().create( - project_id=project_id, - project_name=project_name, - name=name, - stage_type=stage_type, - pause=pause, - prioritized_rulesets=prioritized_rulesets, - description=description, - ) - - -def get_protect_stage( - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, -) -> StageDB | None: - """Retrieves a stage by its ID or name, within a given project. - - Parameters - ---------- - project_id - ID of the project. - project_name - Name of the project. If `project_id` is not provided, - this will be used to look up the project. - stage_id - ID of the stage to retrieve. - stage_name - Name of the stage to retrieve. If `stage_id` is not provided, - this will be used (in conjunction with project ID/name). - - Returns - ------- - The fetched Stage. - """ - return Stages().get(project_id=project_id, project_name=project_name, stage_id=stage_id, stage_name=stage_name) - - -def update_protect_stage( - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, - prioritized_rulesets: Sequence[Ruleset] | None = None, -) -> StageDB | None: - """Updates a stage's rulesets, creating a new version. - - Parameters - ---------- - project_id - ID of the project. - project_name - Name of the project. If `project_id` is not provided, - this will be used to look up the project. - stage_id - ID of the stage to update. - stage_name - Name of the stage to update. If `stage_id` is not provided, - this will be used (in conjunction with project ID/name). - prioritized_rulesets - New list of prioritized rulesets for the stage. - If `None`, effectively clears existing rulesets. - - Returns - ------- - The updated Stage. - """ - return Stages().update( - project_id=project_id, - project_name=project_name, - stage_id=stage_id, - stage_name=stage_name, - prioritized_rulesets=prioritized_rulesets, - ) - - -def pause_protect_stage( - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, -) -> StageDB | None: - """Pauses the specified stage. - - Pauses a stage using either its ID or name within the context of a project. - - Parameters - ---------- - project_id - ID of the project containing the stage. - project_name - Name of the project containing the stage. - (Used if `project_id` is not provided). - stage_id - ID of the stage to pause. - stage_name - Name of the stage to pause. (Used if `stage_id` is not provided). - - Returns - ------- - The Stage with its updated pause state. - """ - return Stages().pause(project_id=project_id, project_name=project_name, stage_id=stage_id, stage_name=stage_name) - - -def resume_protect_stage( - project_id: str | UUID4 | None = None, - project_name: str | None = None, - stage_id: str | UUID4 | None = None, - stage_name: str | None = None, -) -> StageDB | None: - """Resumes a previously paused stage. - - Resumes a stage using either its ID or name within the context of a project. - - Parameters - ---------- - project_id - ID of the project containing the stage. - project_name - Name of the project containing the stage. - (Used if `project_id` is not provided). - stage_id - ID of the stage to resume. - stage_name - Name of the stage to resume. (Used if `stage_id` is not provided). - - Returns - ------- - The Stage with its updated pause state. - """ - return Stages().resume(project_id=project_id, project_name=project_name, stage_id=stage_id, stage_name=stage_name) diff --git a/tests/conftest.py b/tests/conftest.py index 5040d2a..7539722 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -58,8 +58,6 @@ from galileo_core.constants.routes import Routes as CoreRoutes # noqa: E402 from galileo_core.schemas.core.user import User # noqa: E402 from galileo_core.schemas.core.user_role import UserRole # noqa: E402 -from galileo_core.schemas.protect.rule import Rule, RuleOperator # noqa: E402 -from galileo_core.schemas.protect.ruleset import Ruleset # noqa: E402 from splunk_ao.collaborator import CollaboratorRole # noqa: E402 from splunk_ao.config import SplunkAOConfig # noqa: E402 from splunk_ao.configuration import _CONFIGURATION_KEYS, Configuration # noqa: E402 @@ -309,32 +307,6 @@ def _capture_factory(logger): return _capture_factory -@pytest.fixture( - params=[ - # Single ruleset with a single rule. - [Ruleset(rules=[Rule(metric="toxicity", operator=RuleOperator.gt, target_value=0.5)])], - # Single ruleset with multiple rules. - [ - Ruleset( - rules=[ - Rule(metric="toxicity", operator=RuleOperator.gt, target_value=0.5), - Rule(metric="tone", operator=RuleOperator.lt, target_value=0.8), - ] - ) - ], - # Single ruleset with an unknown metric. - [Ruleset(rules=[Rule(metric="unknown", operator=RuleOperator.gt, target_value=0.5)])], - # Multiple rulesets with a single rule each. - [ - Ruleset(rules=[Rule(metric="toxicity", operator=RuleOperator.gt, target_value=0.5)]), - Ruleset(rules=[Rule(metric="toxicity", operator=RuleOperator.lt, target_value=0.8)]), - ], - ] -) -def rulesets(request: pytest.FixtureRequest) -> list[Ruleset]: - return request.param - - @pytest.fixture def enable_galileo_logging(): """Temporarily enable SDK logging for tests that need to capture log output.""" diff --git a/tests/test_logger_batch.py b/tests/test_logger_batch.py index 65f5919..87f624e 100644 --- a/tests/test_logger_batch.py +++ b/tests/test_logger_batch.py @@ -1,7 +1,5 @@ import datetime -import json import logging -import uuid from collections import deque from unittest.mock import AsyncMock, Mock, patch from uuid import UUID, uuid4 @@ -21,9 +19,6 @@ ) from galileo_core.schemas.logging.step import Metrics from galileo_core.schemas.logging.trace import Trace -from galileo_core.schemas.protect.execution_status import ExecutionStatus -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.response import Response, TraceMetadata from galileo_core.schemas.shared.document import Document from galileo_core.schemas.shared.multimodal import ContentModality from splunk_ao.logger import SplunkAOLogger @@ -186,33 +181,6 @@ def test_all_span_types_with_redacted_fields( status_code=200, ) - received_at = int(created_at.timestamp() * 1_000_000_000) - response_at = int((created_at + datetime.timedelta(seconds=1)).timestamp() * 1_000_000_000) - execution_time = 1000.0 - trace_metadata_id = uuid.uuid4() - - logger.add_protect_span( - payload=Payload(input="Protect input", output="Protect output"), - redacted_payload=Payload(input="Protect redacted input", output="Protect redacted output"), - response=Response( - status=ExecutionStatus.triggered, - text="Protect text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - redacted_response=Response( - status=ExecutionStatus.triggered, - text="Protect redacted text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - created_at=created_at, - metadata=metadata, - status_code=200, - ) - logger.add_retriever_span( input="Retriever query with PII: john.doe@email.com", output=["Document with SSN: 123-45-6789", "Document with phone: 555-1234"], @@ -268,36 +236,7 @@ def test_all_span_types_with_redacted_fields( assert tool_span.output == "Tool output with result: result_secret" assert tool_span.redacted_output == "Tool output with result: [REDACTED]" - protect_span = workflow_span.spans[2] - assert isinstance(protect_span, ToolSpan) - assert protect_span.name == "GalileoProtect" - assert json.loads(protect_span.input) == {"input": "Protect input", "output": "Protect output"} - assert json.loads(protect_span.redacted_input) == { - "input": "Protect redacted input", - "output": "Protect redacted output", - } - assert json.loads(protect_span.output) == { - "status": "TRIGGERED", - "text": "Protect text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert json.loads(protect_span.redacted_output) == { - "status": "TRIGGERED", - "text": "Protect redacted text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - - retriever_span = workflow_span.spans[3] + retriever_span = workflow_span.spans[2] assert isinstance(retriever_span, RetrieverSpan) assert retriever_span.input == "Retriever query with PII: john.doe@email.com" assert retriever_span.redacted_input == "Retriever query with PII: [REDACTED]" @@ -430,89 +369,6 @@ def test_add_agent_span(mock_traces_client: Mock, mock_projects_client: Mock, mo assert logger._parent_stack == deque() -@patch("splunk_ao.logger.logger.LogStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_add_protect_tool_span( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock -) -> None: - mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - created_at = datetime.datetime.now() - metadata = {"key": "value"} - logger = SplunkAOLogger(project="my_project", log_stream="my_log_stream") - trace = logger.start_trace( - input="input", name="test-trace", duration_ns=1_000_000, created_at=created_at, metadata=metadata - ) - - received_at = int(created_at.timestamp() * 1_000_000_000) - response_at = int((created_at + datetime.timedelta(seconds=1)).timestamp() * 1_000_000_000) - execution_time = 1000.0 - trace_metadata_id = uuid.uuid4() - - logger.add_protect_span( - payload=Payload(input="Protect input", output="Protect output"), - redacted_payload=Payload(input="Protect redacted input", output="Protect redacted output"), - response=Response( - status=ExecutionStatus.not_triggered, - text="Protect text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - redacted_response=Response( - status=ExecutionStatus.not_triggered, - text="Protect redacted text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - created_at=created_at, - metadata=metadata, - status_code=200, - ) - - logger.conclude(output="response", duration_ns=1_000_000, status_code=200) - logger.flush() - - mock_traces_client_instance.ingest_traces.assert_called_once() - payload = mock_traces_client_instance.ingest_traces.call_args.args[0] - expected_payload = TracesIngestRequest(log_stream_id=None, experiment_id=None, traces=[trace]) - assert payload == expected_payload - protect_span = payload.traces[0].spans[0] - assert isinstance(protect_span, ToolSpan) - assert protect_span.name == "GalileoProtect" - assert json.loads(protect_span.input) == {"input": "Protect input", "output": "Protect output"} - assert json.loads(protect_span.redacted_input) == { - "input": "Protect redacted input", - "output": "Protect redacted output", - } - assert json.loads(protect_span.output) == { - "status": "NOT_TRIGGERED", - "text": "Protect text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert json.loads(protect_span.redacted_output) == { - "status": "NOT_TRIGGERED", - "text": "Protect redacted text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert logger.traces == [] - assert logger._parent_stack == deque() - - @patch("splunk_ao.logger.logger.LogStreams") @patch("splunk_ao.logger.logger.Projects") @patch("splunk_ao.logger.logger.Traces") diff --git a/tests/test_logger_distributed.py b/tests/test_logger_distributed.py index ab5cd76..df479e8 100644 --- a/tests/test_logger_distributed.py +++ b/tests/test_logger_distributed.py @@ -9,9 +9,6 @@ import pytest from galileo_core.schemas.logging.llm import Message -from galileo_core.schemas.protect.execution_status import ExecutionStatus -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.response import Response, TraceMetadata from galileo_core.schemas.shared.document import Document from splunk_ao.logger import SplunkAOLogger from splunk_ao.logger.logger import SplunkAOLoggerException @@ -212,111 +209,6 @@ def test_add_llm_span(mock_traces_client: Mock, mock_projects_client: Mock, mock assert request.spans[0].step_number == 1 -@patch("splunk_ao.logger.logger.LogStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_add_protect_tool_span( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock -) -> None: - mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - created_at = datetime.datetime.now() - metadata = {"key": "value"} - logger = SplunkAOLogger(project="my_project", log_stream="my_log_stream", mode="distributed") - - capture = setup_thread_pool_request_capture(logger) - - logger.start_trace( - input="input", name="test-trace", duration_ns=1_000_000, created_at=created_at, metadata=metadata - ) - - received_at = int(created_at.timestamp() * 1_000_000_000) - response_at = int((created_at + datetime.timedelta(seconds=1)).timestamp() * 1_000_000_000) - execution_time = 1000.0 - trace_metadata_id = uuid.uuid4() - - logger.add_protect_span( - payload=Payload(input="Protect input", output="Protect output"), - redacted_payload=Payload(input="Protect redacted input", output="Protect redacted output"), - response=Response( - status=ExecutionStatus.not_triggered, - text="Protect text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - redacted_response=Response( - status=ExecutionStatus.not_triggered, - text="Protect redacted text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - created_at=created_at, - metadata=metadata, - status_code=200, - ) - - assert len(logger._parent_stack) == 1 - - captured_task = capture.get_task_by_function_name("ingest_traces_with_backoff") - request = captured_task.request - assert isinstance(request, TracesIngestRequest) - - asyncio.run(captured_task.task_func()) - mock_traces_client_instance.ingest_traces.assert_called_with(request) - - assert request.traces[0].input == "input" - assert request.traces[0].name == "test-trace" - assert request.traces[0].created_at == created_at - assert request.traces[0].user_metadata == metadata - assert len(request.traces[0].spans) == 0 - assert request.traces[0].metrics.duration_ns == 1_000_000 - trace_id = request.traces[0].id - - captured_task = capture.get_task_by_function_name("ingest_spans_with_backoff") - request = captured_task.request - assert isinstance(request, SpansIngestRequest) - - asyncio.run(captured_task.task_func()) - mock_traces_client_instance.ingest_spans.assert_called_with(request) - - assert request.trace_id == trace_id - assert request.parent_id == trace_id - protect_span = request.spans[0] - assert protect_span.type == "tool" - assert json.loads(protect_span.input) == {"input": "Protect input", "output": "Protect output"} - assert json.loads(protect_span.redacted_input) == { - "input": "Protect redacted input", - "output": "Protect redacted output", - } - assert json.loads(protect_span.output) == { - "status": "NOT_TRIGGERED", - "text": "Protect text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert json.loads(protect_span.redacted_output) == { - "status": "NOT_TRIGGERED", - "text": "Protect redacted text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert protect_span.name == "GalileoProtect" - assert protect_span.created_at == created_at - assert protect_span.user_metadata == metadata - - @patch("splunk_ao.logger.logger.LogStreams") @patch("splunk_ao.logger.logger.Projects") @patch("splunk_ao.logger.logger.Traces") @@ -1028,25 +920,6 @@ def test_trace_with_multiple_nested_spans( metadata=metadata, ) - received_at = int(created_at.timestamp() * 1_000_000_000) - response_at = int((created_at + datetime.timedelta(seconds=1)).timestamp() * 1_000_000_000) - execution_time = 1000.0 - trace_metadata_id = uuid.uuid4() - - logger.add_protect_span( - payload=Payload(input="Protect input", output="Protect output"), - response=Response( - status=ExecutionStatus.not_triggered, - text="Protect text", - trace_metadata=TraceMetadata( - id=trace_metadata_id, received_at=received_at, response_at=response_at, execution_time=execution_time - ), - ), - created_at=created_at, - metadata=metadata, - status_code=200, - ) - logger.conclude(output="response2", status_code=200, duration_ns=1_000_000) logger.conclude(output="response2", status_code=200, duration_ns=1_000_000) @@ -1054,7 +927,7 @@ def test_trace_with_multiple_nested_spans( assert len(logger._parent_stack) == 0 captured_tasks = capture.get_all_tasks() - assert len(captured_tasks) == 10 + assert len(captured_tasks) == 9 captured_task = captured_tasks[0] assert captured_task.function_name == "ingest_traces_with_backoff" @@ -1186,33 +1059,6 @@ def test_trace_with_multiple_nested_spans( assert request.spans[0].metrics.duration_ns == 1_000_000 captured_task = captured_tasks[7] - assert captured_task.function_name == "ingest_spans_with_backoff" - request = captured_task.request - assert isinstance(request, SpansIngestRequest) - - asyncio.run(captured_task.task_func()) - mock_traces_client_instance.ingest_spans.assert_called_with(request) - - assert request.trace_id == trace_id - assert request.parent_id == workflow_span_id - protect_span = request.spans[0] - assert protect_span.type == "tool" - assert json.loads(protect_span.input) == {"input": "Protect input", "output": "Protect output"} - assert json.loads(protect_span.output) == { - "status": "NOT_TRIGGERED", - "text": "Protect text", - "trace_metadata": { - "id": str(trace_metadata_id), - "received_at": received_at, - "response_at": response_at, - "execution_time": execution_time, - }, - } - assert protect_span.name == "GalileoProtect" - assert protect_span.created_at == created_at - assert protect_span.user_metadata == metadata - - captured_task = captured_tasks[8] assert captured_task.function_name == "update_span_with_backoff" request = captured_task.request assert isinstance(request, SpanUpdateRequest) @@ -1224,7 +1070,7 @@ def test_trace_with_multiple_nested_spans( assert request.output == "response2" assert request.status_code == 200 - captured_task = captured_tasks[9] + captured_task = captured_tasks[8] assert captured_task.function_name == "update_trace_with_backoff" request = captured_task.request assert isinstance(request, TraceUpdateRequest) diff --git a/tests/test_protect.py b/tests/test_protect.py deleted file mode 100644 index c6e7e52..0000000 --- a/tests/test_protect.py +++ /dev/null @@ -1,525 +0,0 @@ -from unittest.mock import ANY, AsyncMock, Mock, patch -from uuid import uuid4 - -from pytest import mark - -from galileo.resources.models.execution_status import ExecutionStatus as APIExecutionStatus -from galileo.resources.models.http_validation_error import HTTPValidationError -from galileo.resources.models.protect_request import ProtectRequest as APIRequest -from galileo.resources.models.protect_response import ProtectResponse as APIResponse -from galileo.resources.models.validation_error import ValidationError -from galileo_core.schemas.protect.execution_status import ExecutionStatus -from galileo_core.schemas.protect.payload import Payload -from galileo_core.schemas.protect.request import Request -from galileo_core.schemas.protect.response import Response -from galileo_core.schemas.protect.rule import Rule, RuleOperator -from galileo_core.schemas.protect.ruleset import Ruleset -from splunk_ao.handlers.langchain.tool import ProtectTool -from splunk_ao.protect import Protect, ainvoke_protect, invoke_protect - -A_PROJECT_NAME = "project_name" -A_STAGE_NAME = "stage_name" -A_PROTECT_INPUT = "invoke" - -# Default values for pass-through parameters. These are used in tests that -# focus on identity-resolution or payload logic, where the specific values -# of timeout/metadata/headers/stage_version don't affect the code path. -DEFAULT_TIMEOUT = 5 -DEFAULT_METADATA: dict | None = None -DEFAULT_HEADERS: dict | None = None -DEFAULT_STAGE_VERSION: int | None = None - - -def invoke_response_data() -> dict: - return { - "action_result": {"action": "PASSTHROUGH"}, - "stage_metadata": {"stage_id": "stage-uuid-example", "stage_name": "test-stage-example"}, - "text": "Original text from payload", - "trace_metadata": { - "id": "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11", - "received_at": 1234567890, - "response_at": 1234567990, - "execution_time": 100.0, - }, - "status": ExecutionStatus.not_triggered.upper(), - "metric_results": {"sentiment_score": 0.95, "toxicity_level": 0.1}, - "ruleset_results": [ - { - "ruleset_id": "ruleset-uuid-example", - "ruleset_name": "basic-ruleset", - "action_result": {"action": "PASSTHROUGH"}, - "rule_results": [ - { - "rule_id": "rule-uuid-example", - "rule_name": "check-sentiment", - "action_result": {"action": "PASSTHROUGH"}, - "metric_results": {"sentiment_score": 0.95}, - "status": ExecutionStatus.not_triggered.upper(), - "conditions_results": [], - } - ], - "status": ExecutionStatus.not_triggered.upper(), - } - ], - "api_version": "1.0.0", - "headers": {"X-Custom-Header": "TestValue"}, - "metadata": {"request_source": "sdk_test"}, - } - - -def invoke_response() -> APIResponse: - return APIResponse.from_dict(invoke_response_data()) - - -# --------------------------------------------------------------------------- -# TestAInvoke: identity-resolution x payload combinations (interacting params) -# Pass-through params use defaults; tested independently below. -# --------------------------------------------------------------------------- -@mark.parametrize( - ("include_project_id", "include_project_name", "include_stage_name", "include_stage_id"), - [ - (True, True, True, True), - (True, False, True, True), - (False, True, True, True), - (False, False, True, True), - (True, True, False, True), - (True, False, False, True), - (True, True, True, False), - (True, False, True, False), - (False, True, False, True), - (False, False, False, True), - ], -) -@mark.parametrize( - "payload", - [ - Payload(input=A_PROTECT_INPUT), - Payload(output=A_PROTECT_INPUT), - Payload(input=A_PROTECT_INPUT, output=A_PROTECT_INPUT), - ], -) -@patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) -class TestAInvoke: - @mark.asyncio - async def test_ainvoke_success( - self, - mock_invoke_post_async: AsyncMock, - include_project_id: bool, - include_project_name: bool, - include_stage_name: bool, - include_stage_id: bool, - payload: Payload, - ) -> None: - mock_invoke_post_async.return_value = invoke_response() - - current_project_id = uuid4() if include_project_id else None - current_project_name = A_PROJECT_NAME if include_project_name else None - current_stage_id = uuid4() if include_stage_id else None - current_stage_name = A_STAGE_NAME if include_stage_name else None - - result = await ainvoke_protect( - payload=payload, - prioritized_rulesets=None, - project_id=current_project_id, - project_name=current_project_name, - stage_id=current_stage_id, - stage_name=current_stage_name, - stage_version=DEFAULT_STAGE_VERSION, - timeout=DEFAULT_TIMEOUT, - metadata=DEFAULT_METADATA, - headers=DEFAULT_HEADERS, - ) - - expected_project_id = str(current_project_id) if current_project_id else None - expected_stage_id = str(current_stage_id) if current_stage_id else None - - body = Request( - payload=payload, - prioritized_rulesets=[], - project_id=expected_project_id, - project_name=current_project_name, - stage_id=expected_stage_id, - stage_name=current_stage_name, - stage_version=DEFAULT_STAGE_VERSION, - timeout=DEFAULT_TIMEOUT, - metadata=DEFAULT_METADATA, - headers=DEFAULT_HEADERS, - ) - - body_dict = body.model_dump(mode="json") - body_dict["prioritized_rulesets"] = body_dict.pop("rulesets", []) - expected_body = APIRequest.from_dict(body_dict) - - mock_invoke_post_async.assert_called_once_with(client=ANY, body=expected_body) - - assert isinstance(result, Response) - - response_data = invoke_response_data() - assert result.text == response_data["text"] - assert result.status == response_data["status"].lower() - assert str(result.trace_metadata.id) == response_data["trace_metadata"]["id"] - - for key in ["action_result", "stage_metadata", "metric_results", "ruleset_results", "api_version"]: - assert key in result.model_extra - assert result.model_extra[key] == response_data[key] - - -# --------------------------------------------------------------------------- -# TestAInvokePassThrough: each pass-through param tested independently -# Uses a single representative identity combo to avoid cross-product explosion. -# --------------------------------------------------------------------------- -class TestAInvokePassThrough: - @mark.parametrize("timeout", [5, 60]) - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_ainvoke_forwards_timeout(self, mock_invoke_post_async: AsyncMock, timeout: float) -> None: - mock_invoke_post_async.return_value = invoke_response() - stage_id = uuid4() - - await ainvoke_protect( - payload=Payload(input=A_PROTECT_INPUT), - prioritized_rulesets=None, - stage_id=stage_id, - stage_name=A_STAGE_NAME, - timeout=timeout, - ) - - body = mock_invoke_post_async.call_args.kwargs["body"] - assert body.timeout == timeout - - @mark.parametrize("metadata", [None, {"key": "value"}]) - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_ainvoke_forwards_metadata(self, mock_invoke_post_async: AsyncMock, metadata: dict | None) -> None: - mock_invoke_post_async.return_value = invoke_response() - stage_id = uuid4() - - await ainvoke_protect( - payload=Payload(input=A_PROTECT_INPUT), - prioritized_rulesets=None, - stage_id=stage_id, - stage_name=A_STAGE_NAME, - metadata=metadata, - ) - - body = mock_invoke_post_async.call_args.kwargs["body"] - if metadata is None: - assert body.metadata is None - else: - assert body.metadata.additional_properties == metadata - - @mark.parametrize("headers", [None, {"key": "value"}]) - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_ainvoke_forwards_headers(self, mock_invoke_post_async: AsyncMock, headers: dict | None) -> None: - mock_invoke_post_async.return_value = invoke_response() - stage_id = uuid4() - - await ainvoke_protect( - payload=Payload(input=A_PROTECT_INPUT), - prioritized_rulesets=None, - stage_id=stage_id, - stage_name=A_STAGE_NAME, - headers=headers, - ) - - body = mock_invoke_post_async.call_args.kwargs["body"] - if headers is None: - assert body.headers is None - else: - assert body.headers.additional_properties == headers - - @mark.parametrize("stage_version", [None, 1, 2]) - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_ainvoke_forwards_stage_version( - self, mock_invoke_post_async: AsyncMock, stage_version: int | None - ) -> None: - mock_invoke_post_async.return_value = invoke_response() - stage_id = uuid4() - - await ainvoke_protect( - payload=Payload(input=A_PROTECT_INPUT), - prioritized_rulesets=None, - stage_id=stage_id, - stage_name=A_STAGE_NAME, - stage_version=stage_version, - ) - - body = mock_invoke_post_async.call_args.kwargs["body"] - assert body.stage_version == stage_version - - -# --------------------------------------------------------------------------- -# TestInvoke: identity-resolution x payload combinations (interacting params) -# Pass-through params use defaults; tested independently below. -# --------------------------------------------------------------------------- -@mark.parametrize( - ("include_project_id", "include_project_name", "include_stage_name", "include_stage_id"), - [ - (True, True, True, True), - (True, False, True, True), - (False, True, True, True), - (False, False, True, True), - (True, True, False, True), - (True, False, False, True), - (True, True, True, False), - (True, False, True, False), - (False, True, False, True), - (False, False, False, True), - ], -) -@mark.parametrize( - "payload", - [ - Payload(input=A_PROTECT_INPUT), - Payload(output=A_PROTECT_INPUT), - Payload(input=A_PROTECT_INPUT, output=A_PROTECT_INPUT), - ], -) -class TestInvoke: - @patch("splunk_ao.protect.ainvoke_protect") - def test_invoke_success( - self, - mock_ainvoke_protect: Mock, - include_project_id: bool, - include_project_name: bool, - include_stage_name: bool, - include_stage_id: bool, - payload: Payload, - ) -> None: - with patch("splunk_ao.protect.async_run") as mock_async_run: - project_id = uuid4() if include_project_id else None - project_name = A_PROJECT_NAME if include_project_name else None - stage_id = uuid4() if include_stage_id else None - stage_name = A_STAGE_NAME if include_stage_name else None - - invoke_protect( - payload=payload, - prioritized_rulesets=None, - project_id=project_id, - project_name=project_name, - stage_id=stage_id, - stage_name=stage_name, - stage_version=DEFAULT_STAGE_VERSION, - timeout=DEFAULT_TIMEOUT, - metadata=DEFAULT_METADATA, - headers=DEFAULT_HEADERS, - ) - mock_async_run.assert_called_once() - mock_ainvoke_protect.assert_called_once_with( - payload=payload, - prioritized_rulesets=None, - project_id=project_id, - project_name=project_name, - stage_id=stage_id, - stage_name=stage_name, - stage_version=DEFAULT_STAGE_VERSION, - timeout=DEFAULT_TIMEOUT, - metadata=DEFAULT_METADATA, - headers=DEFAULT_HEADERS, - ) - - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_langchain_tool( - self, - mock_invoke_post_async: AsyncMock, - include_project_id: bool, - include_project_name: bool, - include_stage_name: bool, - include_stage_id: bool, - payload: Payload, - rulesets: list[Ruleset], - ) -> None: - mock_invoke_post_async.return_value = invoke_response() - - current_project_id = uuid4() if include_project_id else None - current_project_name = A_PROJECT_NAME if include_project_name else None - current_stage_id = uuid4() if include_stage_id else None - current_stage_name = A_STAGE_NAME if include_stage_name else None - - tool = ProtectTool( - prioritized_rulesets=rulesets, - project_id=current_project_id, - project_name=current_project_name, - stage_id=current_stage_id, - stage_name=current_stage_name, - timeout=DEFAULT_TIMEOUT, - stage_version=DEFAULT_STAGE_VERSION, - ) - response_json = tool.run(payload.model_dump()) - assert isinstance(response_json, str) - response = Response.model_validate_json(response_json) - assert response.text is not None - assert response.status == ExecutionStatus.not_triggered - assert mock_invoke_post_async.call_count == 1 - - response_json = await tool.arun(payload.model_dump()) - assert isinstance(response_json, str) - response = Response.model_validate_json(response_json) - assert response.text is not None - assert response.status == ExecutionStatus.not_triggered - assert mock_invoke_post_async.call_count == 2 - - -# --------------------------------------------------------------------------- -# TestInvokePassThrough: each pass-through param tested independently. -# Note: ProtectTool intentionally does NOT support `metadata` and `headers` -# kwargs because those names conflict with the langchain_core tool interface, -# so only `timeout` and `stage_version` are exercised against the tool here. -# --------------------------------------------------------------------------- -class TestInvokePassThrough: - @mark.parametrize("timeout", [5, 60]) - @patch("splunk_ao.protect.ainvoke_protect") - def test_invoke_forwards_timeout(self, mock_ainvoke_protect: Mock, timeout: float) -> None: - with patch("splunk_ao.protect.async_run"): - invoke_protect( - payload=Payload(input=A_PROTECT_INPUT), stage_id=uuid4(), stage_name=A_STAGE_NAME, timeout=timeout - ) - assert mock_ainvoke_protect.call_args.kwargs["timeout"] == timeout - - @mark.parametrize("metadata", [None, {"key": "value"}]) - @patch("splunk_ao.protect.ainvoke_protect") - def test_invoke_forwards_metadata(self, mock_ainvoke_protect: Mock, metadata: dict | None) -> None: - with patch("splunk_ao.protect.async_run"): - invoke_protect( - payload=Payload(input=A_PROTECT_INPUT), stage_id=uuid4(), stage_name=A_STAGE_NAME, metadata=metadata - ) - assert mock_ainvoke_protect.call_args.kwargs["metadata"] == metadata - - @mark.parametrize("headers", [None, {"key": "value"}]) - @patch("splunk_ao.protect.ainvoke_protect") - def test_invoke_forwards_headers(self, mock_ainvoke_protect: Mock, headers: dict | None) -> None: - with patch("splunk_ao.protect.async_run"): - invoke_protect( - payload=Payload(input=A_PROTECT_INPUT), stage_id=uuid4(), stage_name=A_STAGE_NAME, headers=headers - ) - assert mock_ainvoke_protect.call_args.kwargs["headers"] == headers - - @mark.parametrize("stage_version", [None, 1, 2]) - @patch("splunk_ao.protect.ainvoke_protect") - def test_invoke_forwards_stage_version(self, mock_ainvoke_protect: Mock, stage_version: int | None) -> None: - with patch("splunk_ao.protect.async_run"): - invoke_protect( - payload=Payload(input=A_PROTECT_INPUT), - stage_id=uuid4(), - stage_name=A_STAGE_NAME, - stage_version=stage_version, - ) - assert mock_ainvoke_protect.call_args.kwargs["stage_version"] == stage_version - - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_langchain_tool_forwards_timeout(self, mock_invoke_post_async: AsyncMock) -> None: - mock_invoke_post_async.return_value = invoke_response() - tool = ProtectTool(prioritized_rulesets=[], stage_id=uuid4(), stage_name=A_STAGE_NAME, timeout=60) - response_json = tool.run(Payload(input=A_PROTECT_INPUT).model_dump()) - assert isinstance(response_json, str) - body = mock_invoke_post_async.call_args.kwargs["body"] - assert body.timeout == 60 - - @mark.parametrize("stage_version", [None, 1, 2]) - @patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) - @mark.asyncio - async def test_langchain_tool_forwards_stage_version( - self, mock_invoke_post_async: AsyncMock, stage_version: int | None - ) -> None: - mock_invoke_post_async.return_value = invoke_response() - tool = ProtectTool( - prioritized_rulesets=[], stage_id=uuid4(), stage_name=A_STAGE_NAME, stage_version=stage_version - ) - response_json = tool.run(Payload(input=A_PROTECT_INPUT).model_dump()) - assert isinstance(response_json, str) - body = mock_invoke_post_async.call_args.kwargs["body"] - assert body.stage_version == stage_version - - -@patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) -@mark.asyncio -async def test_invoke_with_rulesets(mock_invoke_post_async: Mock) -> None: - mock_invoke_post_async.return_value = invoke_response() - - rules = [Rule(metric="m1", operator=RuleOperator.eq, target_value="v1")] - rulesets = [Ruleset(rules=rules)] - - await ainvoke_protect(payload=Payload(input="test input"), prioritized_rulesets=rulesets, stage_id=uuid4()) - - mock_invoke_post_async.assert_called_once() - api_call_args = mock_invoke_post_async.call_args.kwargs - body = api_call_args["body"] - - assert len(body.prioritized_rulesets) == 1 - api_ruleset = body.prioritized_rulesets[0] - - for i, api_rule in enumerate(api_ruleset.rules): - expected_rule = rules[i] - assert api_rule.metric == expected_rule.metric - assert api_rule.operator.lower() == expected_rule.operator.value.lower() - assert api_rule.target_value == expected_rule.target_value - assert "rulesets" not in body.additional_properties - - -@patch("splunk_ao.protect.invoke_protect_invoke_post.asyncio", new_callable=AsyncMock) -@mark.asyncio -async def test_invoke_api_validation_error(mock_invoke_post_async: Mock) -> None: - error_detail_item = {"loc": ["body", "payload", "input"], "msg": "Field required", "type": "missing"} - validation_error = HTTPValidationError(detail=[error_detail_item]) - mock_invoke_post_async.return_value = validation_error - - payload = Payload(input="Test input") - - result = await ainvoke_protect(payload=payload, stage_id=uuid4()) - assert isinstance(result, HTTPValidationError) - assert result.detail[0]["msg"] == "Field required" - - mock_invoke_post_async.reset_mock() - mock_invoke_post_async.return_value = validation_error - - class_method_result = await Protect().ainvoke(payload=payload, stage_id=uuid4()) - assert isinstance(class_method_result, HTTPValidationError) - assert class_method_result.detail[0]["msg"] == "Field required" - - -@mark.parametrize( - ("status", "expected"), - [ - ("not_triggered", "not_triggered"), - # ("NOT_TRIGGERED", "not_triggered"), - ("Not_Triggered", "not_triggered"), - ("NoT_TRiGgErEd", "not_triggered"), - ("triggered", "triggered"), - ("TRIGGERED", "triggered"), - ], -) -def test_execution_status_case_insensitive(status: str, expected: str) -> None: - """ - Tests that APIExecutionStatus can be initialized - with strings of various casings. - """ - parsed_status = APIExecutionStatus(status) - assert parsed_status.value == expected - - -def test_string_http_validation_error() -> None: - """ - Tests that HTTPValidationError correctly parses a 'detail' field - that is a simple string. - """ - error_message = "A simple error message from API." - error_dict = { - "detail": [{"msg": error_message, "type": "string", "loc": ["api"]}], - "additional_properties": {"key": "value"}, - } - - validation_error = HTTPValidationError.from_dict(error_dict) - - assert isinstance(validation_error, HTTPValidationError) - assert len(validation_error.detail) == 1 - assert len(validation_error.additional_properties) == 1 - detail_item = validation_error.detail[0] - - assert isinstance(detail_item, ValidationError) - assert detail_item.loc == ["api"] - assert detail_item.msg == error_message - assert detail_item.type_ == "string" diff --git a/tests/test_protect_parser.py b/tests/test_protect_parser.py deleted file mode 100644 index db9a6f3..0000000 --- a/tests/test_protect_parser.py +++ /dev/null @@ -1,91 +0,0 @@ -import logging -from json import dumps -from typing import Any -from unittest.mock import patch - -from langchain_core.callbacks import CallbackManagerForLLMRun -from langchain_core.language_models.llms import LLM -from pytest import LogCaptureFixture, mark - -from splunk_ao.handlers.langchain.tool import ProtectParser - -A_TRACE_METADATA_DICT = { - "trace_metadata": { - "id": "57f7ec49-8e44-42cb-8825-4a971e44b252", - "received_at": 1717538501568372000, - "response_at": 1717538501568372000, - "execution_time": 0.46, - } -} - - -class ProtectLLM(LLM): - @property - def _llm_type(self) -> str: - return "protect" - - def _call( - self, - prompt: str, - stop: list[str] | None = None, - run_manager: CallbackManagerForLLMRun | None = None, - **kwargs: Any, - ) -> str: - return prompt - - -@mark.parametrize( - ("output", "ignore_trigger", "expected_return", "expected_call_count"), - [ - (dumps({"text": "foo", "status": "not_triggered", **A_TRACE_METADATA_DICT}), False, "foo", 1), - (dumps({"text": "foo", "status": "not_triggered", **A_TRACE_METADATA_DICT}), True, "foo", 1), - (dumps({"text": "timeout", "status": "TIMEOUT", **A_TRACE_METADATA_DICT}), False, "timeout", 1), - (dumps({"text": "timeout", "status": "TIMEOUT", **A_TRACE_METADATA_DICT}), True, "timeout", 1), - ( - dumps({"text": "not_triggered", "status": "not_triggered", **A_TRACE_METADATA_DICT}), - False, - "not_triggered", - 1, - ), - ( - dumps({"text": "not_triggered", "status": "not_triggered", **A_TRACE_METADATA_DICT}), - True, - "not_triggered", - 1, - ), - ( - dumps({"text": "triggering text", "status": "TRIGGERED", **A_TRACE_METADATA_DICT}), - False, - "triggering text", - 0, - ), - ( - dumps({"text": "triggering text", "status": "TRIGGERED", **A_TRACE_METADATA_DICT}), - True, - "triggering text", - 1, - ), - ], -) -def test_parser(output: str, ignore_trigger: bool, expected_return: str, expected_call_count: int) -> None: - # Verify that the ProtectParser invokes the ProtectLLM only if the status is not "TRIGGERED" - # and ignore_trigger is False. - parser = ProtectParser(chain=ProtectLLM(), ignore_trigger=ignore_trigger) - with patch.object(ProtectLLM, "invoke", wraps=parser.chain.invoke) as mock_fn: - return_value = parser.parser(output) - assert return_value == expected_return - assert mock_fn.call_count == expected_call_count - - -@mark.parametrize(("echo_output", "should_log"), [(True, True), (False, False)]) -def test_echo(echo_output: bool, should_log: bool, caplog: LogCaptureFixture, enable_galileo_logging) -> None: - """Verify that the ProtectParser echoes the output if echo_output is True.""" - parser = ProtectParser(chain=ProtectLLM(), echo_output=echo_output) - - with caplog.at_level(logging.DEBUG, logger="splunk_ao.handlers.langchain.tool"): - parser.parser(dumps({"text": "foo", "status": "not_triggered", **A_TRACE_METADATA_DICT})) - - if should_log: - assert "> Raw response: foo" in caplog.text - else: - assert "> Raw response: foo" not in caplog.text diff --git a/tests/test_stages.py b/tests/test_stages.py deleted file mode 100644 index eeb2c5b..0000000 --- a/tests/test_stages.py +++ /dev/null @@ -1,342 +0,0 @@ -import uuid -from unittest.mock import ANY, Mock, patch - -import pytest -from pydantic import UUID4 - -from galileo.resources.models import HTTPValidationError -from galileo.resources.models.stage_db import StageDB as APIStageDB -from galileo_core.schemas.protect.rule import Rule, RuleOperator -from galileo_core.schemas.protect.ruleset import Ruleset -from galileo_core.schemas.protect.stage import StageDB, StageType -from splunk_ao.stages import ( - create_protect_stage, - get_protect_stage, - pause_protect_stage, - resume_protect_stage, - update_protect_stage, -) - -FIXED_PROJECT_ID = uuid.uuid4() -FIXED_STAGE_ID = uuid.uuid4() - - -def _api_stage_db_factory( - *, - project_id: UUID4 = FIXED_PROJECT_ID, - stage_id: UUID4 | None = None, - name: str = "fixture-stage", - paused: bool = False, - version: int = 1, - stage_type: StageType = StageType.local, -) -> APIStageDB: - data = { - "id": str(stage_id or uuid.uuid4()), - "name": name, - "project_id": str(project_id), - "type": stage_type.value, - "paused": paused, - "version": version, - "created_by": str(uuid.uuid4()), - } - return APIStageDB.from_dict(data) - - -def _core_stage_db_factory( - *, - project_id: UUID4 = FIXED_PROJECT_ID, - stage_id: UUID4 | None = None, - name: str = "fixture-stage", - paused: bool = False, - version: int = 1, - stage_type: StageType = StageType.local, -) -> StageDB: - data = { - "id": str(stage_id or uuid.uuid4()), - "name": name, - "project_id": str(project_id), - "type": stage_type.value, - "paused": paused, - "version": version, - "created_by": str(uuid.uuid4()), - } - return StageDB.model_validate(data) - - -@pytest.fixture(autouse=True) -def _patch_common_modules(): - """Patch common external deps once for the whole module.""" - with patch("splunk_ao.stages.Projects") as proj_patch: - proj_patch.return_value.get_with_env_fallbacks.return_value.id = str(FIXED_PROJECT_ID) - yield - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -@patch("splunk_ao.stages.ts_name", return_value="auto-name") -def test_create_stage_happy_path(mock_ts_name: Mock, mock_api: Mock) -> None: - """Smoke-test: minimal args produce StageDB and correct API call.""" - mock_api.return_value = _api_stage_db_factory(name="auto-name") - - stage = create_protect_stage(project_id=FIXED_PROJECT_ID, pause=False) - - mock_api.assert_called_once_with(project_id=str(FIXED_PROJECT_ID), body=ANY, client=ANY) - assert isinstance(stage, StageDB) - assert stage.project_id == FIXED_PROJECT_ID - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -def test_create_stage_validation_error(mock_api: Mock) -> None: - """create_stage returns HTTPValidationError untouched.""" - err = HTTPValidationError(detail=[{"msg": "bad", "loc": ["body"], "type": "value_error"}]) - mock_api.return_value = err - - res = create_protect_stage(project_id=FIXED_PROJECT_ID, name="x") - assert res is err - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -@patch("splunk_ao.stages.ts_name", return_value="ts_auto") -def test_create_stage_generates_name_and_type(mock_ts_name: Mock, mock_api: Mock) -> None: - """No name provided → ts_name used; stage_type override respected.""" - mock_api.return_value = _api_stage_db_factory(name="ts_auto", stage_type=StageType.central, paused=True) - - stage = create_protect_stage(project_id=FIXED_PROJECT_ID, stage_type=StageType.central, pause=True) - - mock_ts_name.assert_called_once_with("stage") - body = mock_api.call_args.kwargs["body"] - assert body.name == "ts_auto" - assert body.type_ == StageType.central - assert stage.type == StageType.central.value - assert stage.paused is True - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -def test_create_central_stage_with_rulesets(mock_api: Mock) -> None: - rules = [Rule(metric="m1", operator=RuleOperator.eq, target_value="v1")] - rulesets = [Ruleset(rules=rules)] - stage_name = "test-central-stage-with-rules" - mock_api.return_value = _api_stage_db_factory(name=stage_name, stage_type=StageType.central) - - stage_response = create_protect_stage( - project_id=FIXED_PROJECT_ID, name=stage_name, prioritized_rulesets=rulesets, stage_type=StageType.central - ) - - mock_api.assert_called_once() - api_call_args = mock_api.call_args.kwargs - body = api_call_args["body"] - - assert stage_response.type == StageType.central - assert body.name == stage_name - assert body.type_ == StageType.central.value - - assert len(body.prioritized_rulesets) == 1 - api_ruleset = body.prioritized_rulesets[0] - for i, api_rule in enumerate(api_ruleset.rules): - rule = rules[i] - assert api_rule.metric == rule.metric - assert api_rule.operator.lower() == rule.operator.value.lower() - assert api_rule.target_value == rule.target_value - assert "rulesets" not in body.additional_properties - - -@patch("splunk_ao.stages.get_stage_projects_project_id_stages_get.sync") -def test_get_stage_by_id(mock_api: Mock) -> None: - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID) - - stage = get_protect_stage(project_id=FIXED_PROJECT_ID, stage_id=FIXED_STAGE_ID) - - mock_api.assert_called_once_with( - project_id=str(FIXED_PROJECT_ID), stage_id=str(FIXED_STAGE_ID), stage_name=ANY, client=ANY - ) - assert isinstance(stage, StageDB) - assert stage.id == FIXED_STAGE_ID - - -@patch("splunk_ao.stages.get_stage_projects_project_id_stages_get.sync") -@patch("splunk_ao.stages.Projects") -def test_get_stage_by_names(mock_projects_cls: Mock, mock_api: Mock) -> None: - proj_inst = Mock() - proj_inst.get_with_env_fallbacks.return_value.id = str(FIXED_PROJECT_ID) - mock_projects_cls.return_value = proj_inst - - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID, name="named-stage") - - stage = get_protect_stage(project_name="proj", stage_name="named-stage") - - proj_inst.get_with_env_fallbacks.assert_called_once_with(id=None, name="proj") - mock_api.assert_called_once_with( - project_id=str(FIXED_PROJECT_ID), stage_name="named-stage", stage_id=ANY, client=ANY - ) - assert stage.name == "named-stage" - - -@patch("splunk_ao.stages.update_stage_projects_project_id_stages_stage_id_post.sync") -@patch("splunk_ao.stages.Stages.get") -def test_update_stage_rulesets(mock_get: Mock, mock_api: Mock) -> None: - """Verify rulesets payload reaches API, version bumps.""" - mock_get.return_value = _core_stage_db_factory(stage_id=FIXED_STAGE_ID) - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID, version=2) - - rules = [Rule(metric="m", operator=RuleOperator.eq, target_value="v")] - rulesets = [Ruleset(rules=rules)] - stage = update_protect_stage(project_id=FIXED_PROJECT_ID, stage_id=FIXED_STAGE_ID, prioritized_rulesets=rulesets) - - api_body = mock_api.call_args.kwargs["body"] - assert len(api_body.prioritized_rulesets) == 1 - - api_rules = api_body.prioritized_rulesets[0].rules - assert api_rules[0].metric == rules[0].metric - assert api_rules[0].target_value == rules[0].target_value - assert api_rules[0].operator.lower() == rules[0].operator - assert stage.version == 2 - - -@patch("splunk_ao.stages.update_stage_projects_project_id_stages_stage_id_post.sync") -@patch("splunk_ao.stages.Stages.get") -@patch("splunk_ao.stages.Projects") -def test_update_stage_by_names(mock_projects_cls: Mock, mock_get: Mock, mock_api: Mock) -> None: - proj_inst = Mock() - proj_inst.get_with_env_fallbacks.return_value.id = str(FIXED_PROJECT_ID) - mock_projects_cls.return_value = proj_inst - - mock_get.return_value = _core_stage_db_factory(stage_id=FIXED_STAGE_ID, name="named-stage") - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID, version=3) - - stage = update_protect_stage(project_name="proj", stage_name="named-stage", prioritized_rulesets=[]) - - proj_inst.get_with_env_fallbacks.assert_called_once_with(id=None, name="proj") - mock_get.assert_called_once_with(project_id=str(FIXED_PROJECT_ID), stage_name="named-stage") - mock_api.assert_called_once() - assert stage.version == 3 - - -@pytest.mark.parametrize(("pause_flag", "api_fn"), [(True, pause_protect_stage), (False, resume_protect_stage)]) -@patch("splunk_ao.stages.pause_stage_projects_project_id_stages_stage_id_put.sync") -@patch("splunk_ao.stages.Stages.get") -def test_pause_and_resume_by_id(mock_get: Mock, mock_api: Mock, pause_flag, api_fn) -> None: - mock_get.return_value = _core_stage_db_factory(stage_id=FIXED_STAGE_ID, paused=not pause_flag) - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID, paused=pause_flag) - - stage = api_fn(project_id=FIXED_PROJECT_ID, stage_id=FIXED_STAGE_ID) - - mock_api.assert_called_once_with( - project_id=str(FIXED_PROJECT_ID), stage_id=str(FIXED_STAGE_ID), pause=pause_flag, client=ANY - ) - assert stage.paused is pause_flag - - -@patch("splunk_ao.stages.pause_stage_projects_project_id_stages_stage_id_put.sync") -@patch("splunk_ao.stages.Stages.get") -@patch("splunk_ao.stages.Projects") -def test_pause_stage_by_names(mock_projects_cls: Mock, mock_get: Mock, mock_api: Mock) -> None: - proj_inst = Mock() - proj_inst.get_with_env_fallbacks.return_value.id = str(FIXED_PROJECT_ID) - mock_projects_cls.return_value = proj_inst - - mock_get.return_value = _core_stage_db_factory(stage_id=FIXED_STAGE_ID, name="named-stage", paused=False) - mock_api.return_value = _api_stage_db_factory(stage_id=FIXED_STAGE_ID, paused=True) - - stage = pause_protect_stage(project_name="proj", stage_name="named-stage") - - proj_inst.get_with_env_fallbacks.assert_called_once_with(id=None, name="proj") - mock_get.assert_called_once_with(project_id=str(FIXED_PROJECT_ID), stage_name="named-stage") - mock_api.assert_called_once_with( - project_id=str(FIXED_PROJECT_ID), stage_id=str(FIXED_STAGE_ID), pause=True, client=ANY - ) - assert stage.paused is True - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -def test_stage_creation_with_project_id_and_project_name_env_var(mock_api: Mock, monkeypatch) -> None: - monkeypatch.setenv("SPLUNK_AO_PROJECT", "proj") - - rules = [Rule(metric="m1", operator=RuleOperator.eq, target_value="v1")] - rulesets = [Ruleset(rules=rules)] - stage_name = "test-central-stage-with-rules" - mock_api.return_value = _api_stage_db_factory(name=stage_name, stage_type=StageType.central) - - # The project_id passed here should override the SPLUNK_AO_PROJECT env var - stage_response = create_protect_stage( - project_id=FIXED_PROJECT_ID, name=stage_name, prioritized_rulesets=rulesets, stage_type=StageType.central - ) - - mock_api.assert_called_once() - api_call_args = mock_api.call_args.kwargs - body = api_call_args["body"] - - assert stage_response.project_id == FIXED_PROJECT_ID - assert body.name == stage_name - assert body.type_ == StageType.central.value - - assert len(body.prioritized_rulesets) == 1 - api_ruleset = body.prioritized_rulesets[0] - for i, api_rule in enumerate(api_ruleset.rules): - rule = rules[i] - assert api_rule.metric == rule.metric - assert api_rule.operator.lower() == rule.operator.value.lower() - assert api_rule.target_value == rule.target_value - assert "rulesets" not in body.additional_properties - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -def test_stage_creation_with_project_id_and_project_id_env_var(mock_api: Mock, monkeypatch) -> None: - monkeypatch.setenv("SPLUNK_AO_PROJECT_ID", str(FIXED_PROJECT_ID)) - - rules = [Rule(metric="m1", operator=RuleOperator.eq, target_value="v1")] - rulesets = [Ruleset(rules=rules)] - stage_name = "test-central-stage-with-rules" - mock_api.return_value = _api_stage_db_factory(name=stage_name, stage_type=StageType.central) - - # The project_id passed here should override the SPLUNK_AO_PROJECT_ID env var - stage_response = create_protect_stage( - project_id=FIXED_PROJECT_ID, name=stage_name, prioritized_rulesets=rulesets, stage_type=StageType.central - ) - - mock_api.assert_called_once() - api_call_args = mock_api.call_args.kwargs - body = api_call_args["body"] - - assert stage_response.project_id == FIXED_PROJECT_ID - assert body.name == stage_name - assert body.type_ == StageType.central.value - - assert len(body.prioritized_rulesets) == 1 - api_ruleset = body.prioritized_rulesets[0] - for i, api_rule in enumerate(api_ruleset.rules): - rule = rules[i] - assert api_rule.metric == rule.metric - assert api_rule.operator.lower() == rule.operator.value.lower() - assert api_rule.target_value == rule.target_value - assert "rulesets" not in body.additional_properties - - -@patch("splunk_ao.stages.create_stage_projects_project_id_stages_post.sync") -def test_stage_creation_with_project_name_and_project_id_env_var(mock_api: Mock, monkeypatch) -> None: - monkeypatch.setenv("SPLUNK_AO_PROJECT_ID", "proj") - - rules = [Rule(metric="m1", operator=RuleOperator.eq, target_value="v1")] - rulesets = [Ruleset(rules=rules)] - stage_name = "test-central-stage-with-rules" - mock_api.return_value = _api_stage_db_factory(name=stage_name, stage_type=StageType.central) - - # The project_name passed here should override the SPLUNK_AO_PROJECT_ID env var - stage_response = create_protect_stage( - project_name="project name", name=stage_name, prioritized_rulesets=rulesets, stage_type=StageType.central - ) - - mock_api.assert_called_once() - api_call_args = mock_api.call_args.kwargs - body = api_call_args["body"] - - assert stage_response.project_id == FIXED_PROJECT_ID - assert body.name == stage_name - assert body.type_ == StageType.central.value - - assert len(body.prioritized_rulesets) == 1 - api_ruleset = body.prioritized_rulesets[0] - for i, api_rule in enumerate(api_ruleset.rules): - rule = rules[i] - assert api_rule.metric == rule.metric - assert api_rule.operator.lower() == rule.operator.value.lower() - assert api_rule.target_value == rule.target_value - assert "rulesets" not in body.additional_properties