From 410c173c33f46d23aa07745842811ee094f535e8 Mon Sep 17 00:00:00 2001
From: Matteo Mortari <matteo.mortari@gmail.com>
Date: Tue, 3 Sep 2024 07:53:01 +0200
Subject: [PATCH] chore: refactor for constants (#12)

* chore: refactor for constants

Signed-off-by: tarilabs <matteo.mortari@gmail.com>

* Update omlmd/helpers.py

Signed-off-by: Matteo Mortari <matteo.mortari@gmail.com>

Co-authored-by: Isabella Basso <isabellabdoamaral@gmail.com>

---------

Signed-off-by: tarilabs <matteo.mortari@gmail.com>
Co-authored-by: Isabella Basso <isabellabdoamaral@gmail.com>
---
 omlmd/constants.py    |  4 ++++
 omlmd/helpers.py      | 16 +++++++++++-----
 tests/test_helpers.py |  3 ++-
 3 files changed, 17 insertions(+), 6 deletions(-)
 create mode 100644 omlmd/constants.py

diff --git a/omlmd/constants.py b/omlmd/constants.py
new file mode 100644
index 0000000..0354377
--- /dev/null
+++ 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
index f174d73..c85ea51 100644
--- 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
index a4b4f9d..399d527 100644
--- 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