store entire response on push

Signed-off-by: Isabella do Amaral <[email protected]>

omlmd/helpers.py

@@ -29,10 +29,8 @@ def download_file(uri: str):
 
 @dataclass
 class Helper:
-    _registry: OMLMDRegistry = (
-        field(  # TODO: this is a bit limiting when used from CLI, to be refactored
-            default_factory=lambda: OMLMDRegistry(insecure=True)
-        )
+    _registry: OMLMDRegistry = field(
+        default_factory=lambda: OMLMDRegistry(insecure=True)
     )
     _listeners: list[Listener] = field(default_factory=list)
 
@@ -90,13 +88,7 @@ def push(
                 manifest_config=manifest_cfg,
                 do_chunked=True,
             )
-            self.notify_listeners(
-                PushEvent(
-                    result.headers["Docker-Content-Digest"],
-                    target,
-                    model_metadata,
-                )
-            )
+            self.notify_listeners(PushEvent(result, target, model_metadata))
             return result
         finally:
             if owns_meta_files:

omlmd/listener.py

@@ -4,6 +4,8 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
+import requests
+
 from .model_metadata import ModelMetadata
 
 
@@ -26,6 +28,13 @@ class Event(ABC):
 
 @dataclass
 class PushEvent(Event):
-    digest: str
+    response: requests.Response
     target: str
     metadata: ModelMetadata
+
+    @property
+    def ok(self) -> bool:
+        return self.response.status_code == 200
+
+    def get_digest(self) -> str:
+        return self.response.headers["Docker-Content-Digest"] if self.ok else ""

pyproject.toml

@@ -31,6 +31,7 @@ model-registry = ">=0.2.9,<0.3.0"
 ruff = "^0.6.1"
 mypy = "^1.11.1"
 types-pyyaml = "^6.0.12.20240808"
+types-requests = "^2.32.0.20241016"
 
 [tool.poetry.scripts]
 omlmd = "omlmd.cli:cli"

tests/test_e2e_model_registry.py

@@ -43,8 +43,8 @@ class ListenerForModelRegistry(Listener):
 
         def update(self, source: Helper, event: Event) -> None:
             if isinstance(event, PushEvent):
-                self.sha = event.digest
-                self.rm = from_oci_to_kfmr(model_registry, event, event.digest)
+                self.sha = event.get_digest()
+                self.rm = from_oci_to_kfmr(model_registry, event, self.sha)
 
     listener = ListenerForModelRegistry()
     omlmd = Helper()