Allow to specify tag and description when publishing a model version

Closes #137 (closed)

tests/test_elements_worker/test_training.py

@@ -35,11 +35,45 @@ def test_create_archive(model_file_dir):
     assert not os.path.exists(zst_archive_path), "Auto removal failed"
 
 
-def test_create_model_version():
+@pytest.mark.parametrize(
+    "content",
+    [
+        (
+            {
+                "hash": {
+                    "id": "fake_model_version_id",
+                    "model_id": "fake_model_id",
+                    "hash": "hash",
+                    "archive_hash": "archive_hash",
+                    "size": "size",
+                    "tag": "tag",
+                    "description": "description",
+                    "s3_url": "http://hehehe.com",
+                    "s3_put_url": "http://hehehe.com",
+                }
+            }
+        ),
+        (
+            {
+                "hash": {
+                    "id": "fake_model_version_id",
+                    "model_id": "fake_model_id",
+                    "hash": "hash",
+                    "archive_hash": "archive_hash",
+                    "size": "size",
+                    "tag": None,
+                    "description": "",
+                    "s3_url": "http://hehehe.com",
+                    "s3_put_url": "http://hehehe.com",
+                }
+            }
+        ),
+    ],
+)
+def test_create_model_version(content):
     """A new model version is returned"""
 
     model_id = "fake_model_id"
-    model_version_id = "fake_model_version_id"
     training = TrainingWorker()
     training.api_client = MockApiClient()
     model_hash = "hash"
@@ -47,22 +81,10 @@ def test_create_model_version():
     size = "30"
     tag = "tag"
     description = "description"
-    model_version_details = {
-        "id": model_version_id,
-        "model_id": model_id,
-        "hash": model_hash,
-        "archive_hash": archive_hash,
-        "size": size,
-        "tag": tag,
-        "description": description,
-        "s3_url": "http://hehehe.com",
-        "s3_put_url": "http://hehehe.com",
-    }
-
     training.api_client.add_response(
         "CreateModelVersion",
         id=model_id,
-        response=model_version_details,
+        response=content,
         body={
             "hash": model_hash,
             "archive_hash": archive_hash,
@@ -75,7 +97,7 @@ def test_create_model_version():
         training.create_model_version(
             model_id, model_hash, size, archive_hash, tag, description
         )
-        == model_version_details
+        == content
     )
 
 
@@ -90,6 +112,24 @@ def test_create_model_version():
                     "hash": "hash",
                     "archive_hash": "archive_hash",
                     "size": "size",
+                    "tag": "tag",
+                    "description": "description",
+                    "s3_url": "http://hehehe.com",
+                    "s3_put_url": "http://hehehe.com",
+                }
+            },
+            400,
+        ),
+        (
+            {
+                "hash": {
+                    "id": "fake_model_version_id",
+                    "model_id": "fake_model_id",
+                    "hash": "hash",
+                    "archive_hash": "archive_hash",
+                    "size": "size",
+                    "tag": None,
+                    "description": "",
                     "s3_url": "http://hehehe.com",
                     "s3_put_url": "http://hehehe.com",
                 }