diff --git a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py index 3ee7ba482..891fa445f 100644 --- a/monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py +++ b/monkey/tests/unit_tests/common/event_queue/test_pypubsub_agent_event_queue.py @@ -1,4 +1,3 @@ -from dataclasses import dataclass from ipaddress import IPv4Address from typing import Callable, FrozenSet, Union from uuid import UUID @@ -13,7 +12,6 @@ EVENT_TAG_1 = "event tag 1" EVENT_TAG_2 = "event tag 2" -@dataclass(frozen=True) class TestEvent1(AbstractAgentEvent): __test__ = False source: UUID = UUID("f811ad00-5a68-4437-bd51-7b5cc1768ad5") @@ -22,7 +20,6 @@ class TestEvent1(AbstractAgentEvent): tags: FrozenSet = frozenset() -@dataclass(frozen=True) class TestEvent2(AbstractAgentEvent): __test__ = False source: UUID = UUID("e810ad01-6b67-9446-fc58-9b8d717653f7") diff --git a/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py b/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py index b4ddd6612..f83c54ba4 100644 --- a/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py +++ b/monkey/tests/unit_tests/common/event_serializers/test_event_serializer_registry.py @@ -1,25 +1,22 @@ -from dataclasses import dataclass, field from unittest.mock import MagicMock import pytest +from pydantic import Field from common.event_serializers import EventSerializerRegistry, IEventSerializer from common.events import AbstractAgentEvent -@dataclass(frozen=True) class SomeEvent(AbstractAgentEvent): - some_param: int = field(default=435) + some_param: int = Field(default=435) -@dataclass(frozen=True) class OtherEvent(AbstractAgentEvent): - other_param: float = field(default=123.456) + other_param: float = Field(default=123.456) -@dataclass(frozen=True) class NoneEvent(AbstractAgentEvent): - none_param: float = field(default=1.0) + none_param: float = Field(default=1.0) SOME_SERIALIZER = MagicMock(spec=IEventSerializer) diff --git a/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py b/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py index e3cb7a561..a6ba4b6a7 100644 --- a/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py +++ b/monkey/tests/unit_tests/common/event_serializers/test_pydantic_event_serializer.py @@ -1,7 +1,8 @@ from abc import ABC -from dataclasses import dataclass, field +from dataclasses import dataclass import pytest +from pydantic import Field from common.base_models import InfectionMonkeyBaseModel from common.event_serializers import IEventSerializer, PydanticEventSerializer @@ -14,9 +15,8 @@ class NotAgentEvent(ABC): other_field: float -@dataclass(frozen=True) class SomeAgentEvent(AbstractAgentEvent): - bogus: int = field(default_factory=int) + bogus: int = Field(default_factory=int) class PydanticEvent(InfectionMonkeyBaseModel): @@ -28,7 +28,9 @@ def pydantic_event_serializer() -> IEventSerializer: return PydanticEventSerializer(PydanticEvent) -@pytest.mark.parametrize("event", [NotAgentEvent(1, 2.0), SomeAgentEvent(2)]) +@pytest.mark.parametrize( + "event", [NotAgentEvent(some_field=1, other_field=2.0), SomeAgentEvent(bogus=2)] +) def test_pydantic_event_serializer__serialize_wrong_type(pydantic_event_serializer, event): with pytest.raises(TypeError): pydantic_event_serializer.serialize(event)