diff --git a/tests/unit/test_asynchronous.py b/tests/unit/test_asynchronous.py index 1d3f4339..9e14576e 100644 --- a/tests/unit/test_asynchronous.py +++ b/tests/unit/test_asynchronous.py @@ -1,3 +1,4 @@ +import datetime from unittest.mock import AsyncMock import pytest as pytest @@ -5,10 +6,12 @@ AggregateMessage, ForgetMessage, InstanceMessage, + MessageType, PostMessage, ProgramMessage, StoreMessage, ) +from aleph_message.models.execution.environment import MachineResources from aleph_message.status import MessageStatus from aleph.sdk.types import StorageEnum @@ -121,3 +124,63 @@ async def test_forget(mock_session_with_post_success): assert mock_session_with_post_success.http_session.post.called_once assert isinstance(forget_message, ForgetMessage) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "message_type, content", + [ + ( + MessageType.aggregate, + { + "content": {"Hello": datetime.datetime.now()}, + "key": "test", + "address": "0x1", + "time": 1.0, + }, + ), + ( + MessageType.aggregate, + { + "content": {"Hello": datetime.date.today()}, + "key": "test", + "address": "0x1", + "time": 1.0, + }, + ), + ( + MessageType.aggregate, + { + "content": {"Hello": datetime.time()}, + "key": "test", + "address": "0x1", + "time": 1.0, + }, + ), + ( + MessageType.aggregate, + { + "content": { + "Hello": MachineResources( + vcpus=1, + memory=1024, + seconds=1, + ) + }, + "key": "test", + "address": "0x1", + "time": 1.0, + }, + ), + ], +) +async def test_prepare_aleph_message( + mock_session_with_post_success, message_type, content +): + # Call the function under test + async with mock_session_with_post_success as session: + await session._prepare_aleph_message( + message_type=message_type, + content=content, + channel="TEST", + )