Allow to specify tag and description when publishing a model version

Closes #137 (closed)

tests/test_elements_worker/test_training.py

@@ -36,55 +36,42 @@ def test_create_archive(model_file_dir):
 
 
 @pytest.mark.parametrize(
-    "content",
+    "tag, description",
     [
-        (
-            {
-                "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",
-                }
-            }
-        ),
+        ("tag", "description"),
+        (None, "description"),
+        ("", "description"),
+        ("tag", ""),
+        ("", ""),
+        (None, ""),
     ],
 )
-def test_create_model_version(content):
+def test_create_model_version(tag, description):
     """A new model version is returned"""
 
+    model_version_id = "fake_model_version_id"
     model_id = "fake_model_id"
     training = TrainingWorker()
     training.api_client = MockApiClient()
     model_hash = "hash"
     archive_hash = "archive_hash"
     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=content,
+        response=model_version_details,
         body={
             "hash": model_hash,
             "archive_hash": archive_hash,
@@ -97,7 +84,7 @@ def test_create_model_version(content):
         training.create_model_version(
             model_id, model_hash, size, archive_hash, tag, description
         )
-        == content
+        == model_version_details
     )