chore: refactor for constants (#12)

* chore: refactor for constants Signed-off-by: tarilabs <[email protected]> * Update omlmd/helpers.py Signed-off-by: Matteo Mortari <[email protected]> Co-authored-by: Isabella Basso <[email protected]> --------- Signed-off-by: tarilabs <[email protected]> Co-authored-by: Isabella Basso <[email protected]>
containers · Sep 3, 2024 · 410c173 · 410c173
1 parent 3c87873
commit 410c173
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/omlmd/constants.py b/omlmd/constants.py
@@ -0,0 +1,4 @@
+FILENAME_METADATA_JSON = "model_metadata.omlmd.json"
+FILENAME_METADATA_YAML = "model_metadata.omlmd.yaml"
+MIME_APPLICATION_CONFIG = "application/x-config"
+MIME_APPLICATION_MLMODEL = "application/x-mlmodel"
diff --git a/omlmd/helpers.py b/omlmd/helpers.py
@@ -7,6 +7,12 @@
 from dataclasses import fields
 from pathlib import Path
 
+from omlmd.constants import (
+    FILENAME_METADATA_JSON,
+    FILENAME_METADATA_YAML,
+    MIME_APPLICATION_CONFIG,
+    MIME_APPLICATION_MLMODEL,
+)
 from omlmd.listener import Event, Listener, PushEvent
 from omlmd.model_metadata import ModelMetadata
 from omlmd.provider import OMLMDRegistry
@@ -65,8 +71,8 @@ def push(
         if isinstance(path, str):
             path = Path(path)
 
-        json_meta = path.parent / "model_metadata.omlmd.json"
-        yaml_meta = path.parent / "model_metadata.omlmd.yaml"
+        json_meta = path.parent / FILENAME_METADATA_JSON
+        yaml_meta = path.parent / FILENAME_METADATA_YAML
         if model_metadata.is_empty() and json_meta.exists() and yaml_meta.exists():
             owns_meta_files = False
             logger.warning("No metadata supplied, but reusing md files found in path.")
@@ -81,11 +87,11 @@ def push(
             json_meta.write_text(model_metadata.to_json())
             yaml_meta.write_text(model_metadata.to_yaml())
 
-        manifest_cfg = f"{json_meta}:application/x-config"
+        manifest_cfg = f"{json_meta}:{MIME_APPLICATION_CONFIG}"
         files = [
-            f"{path}:application/x-mlmodel",
+            f"{path}:{MIME_APPLICATION_MLMODEL}",
             manifest_cfg,
-            f"{yaml_meta}:application/x-config",
+            f"{yaml_meta}:{MIME_APPLICATION_CONFIG}",
         ]
         try:
             # print(target, files, model_metadata.to_annotations_dict())

diff --git a/tests/test_helpers.py b/tests/test_helpers.py
@@ -4,6 +4,7 @@
 
 import pytest
 
+from omlmd.constants import MIME_APPLICATION_MLMODEL
 from omlmd.helpers import Helper
 from omlmd.listener import Event, Listener
 from omlmd.model_metadata import ModelMetadata, deserialize_mdfile
@@ -89,5 +90,5 @@ def test_e2e_push_pull_with_filters(tmp_path, target):
         author="John Doe",
         accuracy=0.987,
     )
-    omlmd.pull(target, tmp_path, media_types=["application/x-mlmodel"])
+    omlmd.pull(target, tmp_path, media_types=[MIME_APPLICATION_MLMODEL])
     assert len(list(tmp_path.iterdir())) == 1