From 556d10905a8a90da19c92c8d200fd257629354cd Mon Sep 17 00:00:00 2001
From: Isabella do Amaral <idoamara@redhat.com>
Date: Fri, 25 Oct 2024 10:29:20 -0300
Subject: [PATCH] store entire response on push

Signed-off-by: Isabella do Amaral <idoamara@redhat.com>
---
 omlmd/helpers.py                 | 14 +++-----------
 omlmd/listener.py                | 11 ++++++++++-
 poetry.lock                      | 16 +++++++++++++++-
 pyproject.toml                   |  1 +
 tests/test_e2e_model_registry.py |  4 ++--
 5 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/omlmd/helpers.py b/omlmd/helpers.py
index faa481e..38ad5e4 100644
--- a/omlmd/helpers.py
+++ b/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:
diff --git a/omlmd/listener.py b/omlmd/listener.py
index 51583fa..7e15b62 100644
--- a/omlmd/listener.py
+++ b/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 ""
diff --git a/poetry.lock b/poetry.lock
index 6bd7e54..31ce5b5 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2397,6 +2397,20 @@ files = [
     {file = "types_PyYAML-6.0.12.20240808-py3-none-any.whl", hash = "sha256:deda34c5c655265fc517b546c902aa6eed2ef8d3e921e4765fe606fe2afe8d35"},
 ]
 
+[[package]]
+name = "types-requests"
+version = "2.32.0.20241016"
+description = "Typing stubs for requests"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "types-requests-2.32.0.20241016.tar.gz", hash = "sha256:0d9cad2f27515d0e3e3da7134a1b6f28fb97129d86b867f24d9c726452634d95"},
+    {file = "types_requests-2.32.0.20241016-py3-none-any.whl", hash = "sha256:4195d62d6d3e043a4eaaf08ff8a62184584d2e8684e9d2aa178c7915a7da3747"},
+]
+
+[package.dependencies]
+urllib3 = ">=2"
+
 [[package]]
 name = "typing-extensions"
 version = "4.12.2"
@@ -2568,4 +2582,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools",
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.9"
-content-hash = "782057572cf1e420e7dae3604e26d3ed32682762eb45ffc83f4b08691958aa04"
+content-hash = "09148e1d60c1b8c92d31467b68e406f6dbf2566f3d0b0cff7f7f01067276d726"
diff --git a/pyproject.toml b/pyproject.toml
index fa0b89f..de386f5 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/tests/test_e2e_model_registry.py b/tests/test_e2e_model_registry.py
index 748e2e1..39adc26 100644
--- a/tests/test_e2e_model_registry.py
+++ b/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()