Add parse_volume to fix parsing of MachineVolume types

Co-authored-by: Hugo Herter <[email protected]>

Reformat omfg

Author: MHHukiewitz

src/aleph/sdk/client/authenticated_http.py

@@ -37,7 +37,7 @@
 from ..conf import settings
 from ..exceptions import BroadcastError, InsufficientFundsError, InvalidMessageError
 from ..types import Account, StorageEnum
-from ..utils import extended_json_encoder
+from ..utils import extended_json_encoder, parse_volume
 from .abstract import AuthenticatedAlephClient
 from .http import AlephHttpClient
 
@@ -449,9 +449,7 @@ async def create_program(
             # Trigger on HTTP calls.
             triggers = {"http": True, "persistent": persistent}
 
-        volumes: List[MachineVolume] = [
-            MachineVolume.parse_obj(volume) for volume in volumes
-        ]
+        volumes: List[MachineVolume] = [parse_volume(volume) for volume in volumes]
 
         content = ProgramContent(
             type="vm-function",
@@ -482,7 +480,7 @@ async def create_program(
                 if runtime == settings.DEFAULT_RUNTIME_ID
                 else "",
             ),
-            volumes=[MachineVolume.parse_obj(volume) for volume in volumes],
+            volumes=[parse_volume(volume) for volume in volumes],
             time=time.time(),
             metadata=metadata,
         )
@@ -555,7 +553,7 @@ async def create_instance(
                 if rootfs == settings.DEFAULT_RUNTIME_ID
                 else "",
             ),
-            volumes=[MachineVolume.parse_obj(volume) for volume in volumes],
+            volumes=[parse_volume(volume) for volume in volumes],
             time=time.time(),
             authorized_keys=ssh_keys,
             metadata=metadata,

src/aleph/sdk/utils.py

@@ -5,11 +5,23 @@
 from enum import Enum
 from pathlib import Path
 from shutil import make_archive
-from typing import Any, Iterable, Optional, Protocol, Tuple, Type, TypeVar, Union
+from typing import (
+    Any,
+    Iterable,
+    Mapping,
+    Optional,
+    Protocol,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+)
 from zipfile import BadZipFile, ZipFile
 
 from aleph_message.models import MessageType
 from aleph_message.models.execution.program import Encoding
+from aleph_message.models.execution.volume import MachineVolume
 from pydantic.json import pydantic_encoder
 
 from aleph.sdk.conf import settings
@@ -150,3 +162,19 @@ def extended_json_encoder(obj: Any) -> Any:
         return obj.hour * 3600 + obj.minute * 60 + obj.second + obj.microsecond / 1e6
     else:
         return pydantic_encoder(obj)
+
+
+def parse_volume(volume_dict: Union[Mapping, MachineVolume]) -> MachineVolume:
+    # Python 3.9 does not support `isinstance(volume_dict, MachineVolume)`,
+    # so we need to iterate over all types.
+    if any(
+        isinstance(volume_dict, volume_type) for volume_type in get_args(MachineVolume)
+    ):
+        return volume_dict
+    for volume_type in get_args(MachineVolume):
+        try:
+            return volume_type.parse_obj(volume_dict)
+        except ValueError:
+            continue
+    else:
+        raise ValueError(f"Could not parse volume: {volume_dict}")

tests/unit/test_utils.py

@@ -12,8 +12,13 @@
     StoreMessage,
 )
 from aleph_message.models.execution.environment import MachineResources
+from aleph_message.models.execution.volume import (
+    EphemeralVolume,
+    ImmutableVolume,
+    PersistentVolume,
+)
 
-from aleph.sdk.utils import enum_as_str, get_message_type_value
+from aleph.sdk.utils import enum_as_str, get_message_type_value, parse_volume
 
 
 def test_get_message_type_value():
@@ -128,3 +133,44 @@ async def test_prepare_aleph_message(
         )
 
     assert message.content.dict() == content
+
+
+def test_parse_immutable_volume():
+    volume_dict = {
+        "ref": "QmX8K1c22WmQBAww5ShWQqwMiFif7XFrJD6iFBj7skQZXW",
+        "use_latest": True,
+        "comment": "Dummy hash",
+    }
+    volume = parse_volume(volume_dict)
+    volume = parse_volume(volume)
+    assert volume
+    assert isinstance(volume, ImmutableVolume)
+
+
+def test_parse_ephemeral_volume():
+    volume_dict = {
+        "comment": "Dummy hash",
+        "ephemeral": True,
+        "size_mib": 1,
+    }
+    volume = parse_volume(volume_dict)
+    volume = parse_volume(volume)
+    assert volume
+    assert isinstance(volume, EphemeralVolume)
+
+
+def test_parse_persistent_volume():
+    volume_dict = {
+        "parent": {
+            "ref": "QmX8K1c22WmQBAww5ShWQqwMiFif7XFrJD6iFBj7skQZXW",
+            "use_latest": True,
+            "comment": "Dummy hash",
+        },
+        "persistence": "host",
+        "name": "test",
+        "size_mib": 1,
+    }
+    volume = parse_volume(volume_dict)
+    volume = parse_volume(volume)
+    assert volume
+    assert isinstance(volume, PersistentVolume)