Skip to content

Commit 364b7ec

Browse files
authored
Solve type saving issues (#134)
* Fix: Solved type saving issue ensuring to get proper attributes on model validations. * Fix: Fixed class method signature to use the good one from Pydantic, they introduced in their last version 2.12.
1 parent b1ab9f1 commit 364b7ec

File tree

2 files changed

+83
-18
lines changed

2 files changed

+83
-18
lines changed

aleph_message/models/execution/instance.py

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from typing import List, Optional
44

55
from pydantic import Field, model_validator
6+
from typing_extensions import Self
67

78
from aleph_message.models.abstract import HashableModel
89

@@ -43,38 +44,37 @@ class InstanceContent(BaseExecutableContent):
4344
)
4445

4546
@model_validator(mode="after")
46-
def check_requirements(cls, values):
47-
if values.requirements:
47+
def check_requirements(self) -> Self:
48+
if self.requirements:
49+
if (
50+
self.payment and (self.payment.is_stream or self.payment.is_credit)
51+
) and (not self.requirements.node or not self.requirements.node.node_hash):
52+
raise ValueError(
53+
"Node hash assignment is needed for PAYG or Credit payments"
54+
)
4855
# GPU filter only supported for QEmu instances with node_hash assigned
49-
if values.requirements.gpu:
50-
if (
51-
not values.requirements.node
52-
or not values.requirements.node.node_hash
53-
):
56+
if self.requirements.gpu:
57+
if not self.requirements.node or not self.requirements.node.node_hash:
5458
raise ValueError("Node hash assignment is needed for GPU support")
5559

5660
if (
57-
values.environment
58-
and values.environment.hypervisor != HypervisorType.qemu
61+
self.environment
62+
and self.environment.hypervisor != HypervisorType.qemu
5963
):
6064
raise ValueError("GPU option is only supported for QEmu hypervisor")
6165

6266
# Terms and conditions filter only supported for PAYG/coco instances with node_hash assigned
63-
if (
64-
values.requirements.node
65-
and values.requirements.node.terms_and_conditions
66-
):
67-
if not values.requirements.node.node_hash:
67+
if self.requirements.node and self.requirements.node.terms_and_conditions:
68+
if not self.requirements.node.node_hash:
6869
raise ValueError(
6970
"Terms_and_conditions field needs a requirements.node.node_hash value"
7071
)
7172

7273
if (
73-
not values.payment.is_stream
74-
and not values.environment.trusted_execution
75-
):
74+
not self.payment or not self.payment.is_stream
75+
) and not self.environment.trusted_execution:
7676
raise ValueError(
7777
"Only PAYG/coco instances can have a terms_and_conditions"
7878
)
7979

80-
return values
80+
return self

aleph_message/tests/test_models.py

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,71 @@ def test_validation_on_gpu_hypervisor_options():
245245
) # Ignore "Value error, ..."
246246

247247

248+
def test_validation_on_payg_payment_options():
249+
"""Ensure that a node_hash option is required for stream payments."""
250+
path = Path(__file__).parent / "messages/instance_gpu_machine.json"
251+
message_dict = json.loads(path.read_text())
252+
message_dict["content"]["requirements"]["gpu"] = None
253+
message = create_new_message(message_dict, factory=InstanceMessage)
254+
255+
assert isinstance(message, InstanceMessage)
256+
assert hash(message.content)
257+
assert message.content.payment
258+
assert message.content.payment.type == PaymentType.superfluid
259+
assert message.content.environment.hypervisor == HypervisorType.qemu
260+
261+
message_dict["content"]["requirements"]["node"] = None
262+
try:
263+
_ = create_new_message(message_dict, factory=InstanceMessage)
264+
raise AssertionError("An exception should have been raised before this point.")
265+
except ValidationError as e:
266+
assert e.errors()[0]["loc"] in [("content",), ("content", "__root__")]
267+
268+
error_msg = e.errors()[0]["msg"]
269+
assert (
270+
"Node hash assignment is needed for PAYG or Credit payments" in error_msg
271+
) # Ignore "Value error, ..."
272+
273+
274+
def test_validation_on_credits_payment_options():
275+
"""Ensure that a node_hash option is required for credit payments."""
276+
path = Path(__file__).parent / "messages/instance_gpu_machine.json"
277+
message_dict = json.loads(path.read_text())
278+
# Patch the gpu field with some info
279+
message_dict["content"]["payment"]["type"] = "credit"
280+
message_dict["content"]["payment"]["receiver"] = None
281+
message = create_new_message(message_dict, factory=InstanceMessage)
282+
283+
assert isinstance(message, InstanceMessage)
284+
assert hash(message.content)
285+
assert message.content.payment
286+
assert message.content.payment.type == PaymentType.credit
287+
assert message.content.environment.hypervisor == HypervisorType.qemu
288+
289+
# Remove gpu field with some info
290+
message_dict["content"]["requirements"]["gpu"] = None
291+
message = create_new_message(message_dict, factory=InstanceMessage)
292+
293+
assert isinstance(message, InstanceMessage)
294+
assert hash(message.content)
295+
assert message.content.payment
296+
assert message.content.payment.type == PaymentType.credit
297+
assert message.content.environment.hypervisor == HypervisorType.qemu
298+
assert message.content.requirements.gpu is None
299+
300+
message_dict["content"]["requirements"]["node"] = None
301+
try:
302+
_ = create_new_message(message_dict, factory=InstanceMessage)
303+
raise AssertionError("An exception should have been raised before this point.")
304+
except ValidationError as e:
305+
assert e.errors()[0]["loc"] in [("content",), ("content", "__root__")]
306+
307+
error_msg = e.errors()[0]["msg"]
308+
assert (
309+
"Node hash assignment is needed for PAYG or Credit payments" in error_msg
310+
) # Ignore "Value error, ..."
311+
312+
248313
def test_message_machine_port_mapping():
249314
message_dict = {
250315
"chain": "ETH",

0 commit comments

Comments
 (0)