Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat+fix(engine)!: Support enums and defaults in Action Templates #644

Merged
merged 33 commits into from
Dec 23, 2024
Merged
Show file tree
Hide file tree
Changes from 31 commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
3fb7eb9
Add tests
topher-lo Dec 19, 2024
fceab7a
Add more informative inline comments in tests
topher-lo Dec 19, 2024
3442d77
Test defaults logic in template actions
topher-lo Dec 20, 2024
32460bf
Add enum type to expectations
topher-lo Dec 20, 2024
ace09b1
fix missing arg in test
topher-lo Dec 20, 2024
3513263
Update verbosity of register logs
topher-lo Dec 21, 2024
29436ad
Update default / non default args tests
topher-lo Dec 21, 2024
e17e75e
feat: Use validated args in `run_template_action`
topher-lo Dec 21, 2024
b165b08
Replace `None` with `null` (reverved string) to represent null defaul…
topher-lo Dec 21, 2024
fc40087
Drop unused validate path in _run_action_direct
topher-lo Dec 21, 2024
9d5495e
Check RegistryValidationError raised
topher-lo Dec 21, 2024
e72fbff
return validated_args in BoundRegistryAction.validate_args with json …
topher-lo Dec 21, 2024
134794a
fix replace str | None with str | null in wazuh `run_rootcheck`
topher-lo Dec 21, 2024
4fc933f
fix int | null in _example_template_action
topher-lo Dec 21, 2024
1bb2bc8
Replace int | None with int | null in test_expectations schemas
topher-lo Dec 21, 2024
b924857
Revert null type
topher-lo Dec 21, 2024
112deb1
revert nulls
topher-lo Dec 21, 2024
a5a1194
revert null : None mapping
topher-lo Dec 21, 2024
794fe98
validate args in run action direct (udf)
topher-lo Dec 21, 2024
4e4ba4c
Apply suggestions from code review
topher-lo Dec 22, 2024
3e1d6aa
Updates
topher-lo Dec 22, 2024
80cc4b8
Support single quotes in enum
topher-lo Dec 22, 2024
d39f06f
Add Enum as prefix to enum class name
topher-lo Dec 22, 2024
7398642
refactor: Call TEMPLATE_ACTION_INPUTS var `validated_args`
topher-lo Dec 22, 2024
c9fa57c
Capture e representation for unexpected errors in action.validate_args
topher-lo Dec 22, 2024
b5d70ee
Drop RegistryValidationError from registry
topher-lo Dec 23, 2024
8091d72
Make RegistryValidationError serializable
topher-lo Dec 23, 2024
67bf6fe
Merge branch 'main' into feat/enum-action-templates
topher-lo Dec 23, 2024
a8e09fa
Fix expected enum class value with prefix
topher-lo Dec 23, 2024
402fe85
more test fixes
topher-lo Dec 23, 2024
15287bf
simplify enum validate schema param names
topher-lo Dec 23, 2024
b7373ce
Update models.py
topher-lo Dec 23, 2024
90f26c4
fix missing raise RegistryValidationError(
topher-lo Dec 23, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 2 additions & 6 deletions registry/tracecat_registry/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,8 @@
"Could not import tracecat. Please install `tracecat` to use the registry."
) from None

from tracecat_registry._internal import registry, secrets
from tracecat_registry._internal.exceptions import ( # noqa: E402
RegistryActionError,
RegistryValidationError,
)
from tracecat_registry._internal import exceptions, registry, secrets
from tracecat_registry._internal.exceptions import RegistryActionError
from tracecat_registry._internal.logger import logger
from tracecat_registry._internal.models import RegistrySecret

Expand All @@ -24,6 +21,5 @@
"logger",
"secrets",
"exceptions",
"RegistryValidationError",
"RegistryActionError",
]
11 changes: 0 additions & 11 deletions registry/tracecat_registry/_internal/exceptions.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
from typing import Any

from pydantic_core import ValidationError


class TracecatException(Exception):
"""Tracecat generic user-facing exception"""
Expand All @@ -13,12 +11,3 @@ def __init__(self, *args, detail: Any | None = None, **kwargs):

class RegistryActionError(TracecatException):
"""Exception raised when a registry UDF error occurs."""


class RegistryValidationError(TracecatException):
"""Exception raised when a registry validation error occurs."""

def __init__(self, *args, key: str, err: ValidationError | str | None = None):
super().__init__(*args)
self.key = key
self.err = err
86 changes: 85 additions & 1 deletion tests/unit/test_expectation.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import json
from datetime import datetime

import lark
import pytest
from pydantic import ValidationError

Expand Down Expand Up @@ -247,4 +248,87 @@ def test_validate_schema_failure():
assert "end_time" in str(e.value)


# ... existing tests ...
@pytest.mark.parametrize(
"status,priority",
[
("PENDING", "low"),
("running", "low"),
("Completed", "low"),
],
)
def test_validate_schema_with_enum(status, priority):
schema = {
"status": {
"type": 'enum["PENDING", "running", "Completed"]',
"description": "The status of the job",
},
"priority": {
"type": 'enum["low", "medium", "high"]',
"description": "The priority level",
"default": "low",
},
}

mapped = {k: ExpectedField(**v) for k, v in schema.items()}
model = create_expectation_model(mapped)

# Test with provided priority
model_instance = model(status=status, priority=priority)
assert model_instance.status.__class__.__name__ == "EnumStatus"
assert model_instance.priority.__class__.__name__ == "EnumPriority"

# Test default priority
model_instance_default = model(status=status)
assert str(model_instance_default.priority) == "low"


@pytest.mark.parametrize(
"schema_def,error_type,error_message",
[
(
{"status": {"type": "enum[]", "description": "Empty enum"}},
lark.exceptions.UnexpectedCharacters,
"No terminal matches ']'",
),
(
{
"status": {
"type": 'enum["Pending", "PENDING"]',
"description": "Duplicate values",
}
},
lark.exceptions.VisitError,
"Duplicate enum value",
),
],
)
def test_validate_schema_with_invalid_enum_definition(
schema_def, error_type, error_message
):
with pytest.raises(error_type, match=error_message):
mapped = {k: ExpectedField(**v) for k, v in schema_def.items()}
create_expectation_model(mapped)


@pytest.mark.parametrize(
"invalid_value",
[
"invalid_status",
"INVALID",
"pending!",
"",
],
)
def test_validate_schema_with_invalid_enum_values(invalid_value):
schema = {
"status": {
"type": 'enum["PENDING", "running", "Completed"]',
"description": "The status of the job",
}
}

mapped = {k: ExpectedField(**v) for k, v in schema.items()}
DynamicModel = create_expectation_model(mapped)

with pytest.raises(ValidationError):
DynamicModel(status=invalid_value)
2 changes: 1 addition & 1 deletion tests/unit/test_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@
import textwrap

import pytest
from tracecat_registry import RegistryValidationError

from tracecat.concurrency import GatheringTaskGroup
from tracecat.registry.actions.models import RegistryActionRead
from tracecat.registry.actions.service import RegistryActionsService
from tracecat.registry.repository import Repository
from tracecat.types.exceptions import RegistryValidationError


@pytest.fixture
Expand Down
Loading
Loading