Skip to content

Commit

Permalink
Make event validation stricter
Browse files Browse the repository at this point in the history
Add extra="forbid" and strict=True to all the events
Add missing ensemble field for EESnapshot and EESnapshotUpdate
  • Loading branch information
JHolba committed Sep 5, 2024
1 parent f60497e commit 595a08f
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 8 deletions.
13 changes: 7 additions & 6 deletions src/_ert/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
else:
from typing import Annotated

from pydantic import BaseModel, Field, TypeAdapter
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter


class Id:
Expand Down Expand Up @@ -65,6 +65,7 @@ class Id:


class BaseEvent(BaseModel):
model_config = ConfigDict(strict=True, extra="forbid")
time: datetime = Field(default_factory=datetime.now)


Expand Down Expand Up @@ -159,12 +160,12 @@ class EnsembleCancelled(EnsembleBaseEvent):
event_type: Id.ENSEMBLE_CANCELLED_TYPE = Id.ENSEMBLE_CANCELLED


class EESnapshot(BaseEvent):
class EESnapshot(EnsembleBaseEvent):
event_type: Id.EE_SNAPSHOT_TYPE = Id.EE_SNAPSHOT
snapshot: Any


class EESnapshotUpdate(BaseEvent):
class EESnapshotUpdate(EnsembleBaseEvent):
event_type: Id.EE_SNAPSHOT_UPDATE_TYPE = Id.EE_SNAPSHOT_UPDATE
snapshot: Any

Expand Down Expand Up @@ -230,15 +231,15 @@ class EEUserDone(BaseEvent):


def dispatch_event_from_json(raw_msg: Union[str, bytes]) -> DispatchEvent:
return DispatchEventAdapter.validate_json(raw_msg, strict=True)
return DispatchEventAdapter.validate_json(raw_msg)


def event_from_json(raw_msg: Union[str, bytes]) -> Event:
return EventAdapter.validate_json(raw_msg, strict=True)
return EventAdapter.validate_json(raw_msg)


def event_from_dict(dict_msg: Dict[str, Any]) -> Event:
return EventAdapter.validate_python(dict_msg, strict=True)
return EventAdapter.validate_python(dict_msg)


def event_to_json(event: Event) -> str:
Expand Down
4 changes: 2 additions & 2 deletions tests/performance_tests/test_snapshot.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,8 @@ def simulate_forward_model_event_handling(
ensemble=ens_id,
real=str(real),
fm_step=str(fm_idx),
stderr="foo",
stdout="bar",
std_err="foo",
std_out="bar",
)
)
for current_memory_usage in range(memory_reports):
Expand Down

0 comments on commit 595a08f

Please sign in to comment.