Addressed PR comments

seanpmorgan · May 15, 2019 · b24e3cc · b24e3cc
1 parent 600d769
commit b24e3cc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/model-archiver/model_archiver/tests/integ_tests/test_integration_model_archiver.py b/model-archiver/model_archiver/tests/integ_tests/test_integration_model_archiver.py
@@ -14,12 +14,12 @@
 
 
 def update_tests(test):
-    test["modelName"] = test.get("modelName") if test.get("modelName") is not None else DEFAULT_MODEL_NAME
-    test["modelPath"] = test.get("modelPath") if test.get("modelPath") is not None else DEFAULT_MODEL_PATH
-    test["handler"] = test.get("handler") if test.get("handler") is not None else DEFAULT_HANDLER
-    test["runtime"] = test.get("runtime") if test.get("runtime") is not None else DEFAULT_RUNTIME
-    test["exportPath"] = test.get("exportPath") if test.get("exportPath") is not None else DEFAULT_EXPORT_PATH
-    test["archiveFormat"] = test.get("archiveFormat") if test.get("archiveFormat") is not None else "default"
+    test["modelName"] = test.get("modelName", DEFAULT_MODEL_NAME)
+    test["modelPath"] = test.get("modelPath", DEFAULT_MODEL_PATH)
+    test["handler"] = test.get("handler", DEFAULT_HANDLER)
+    test["runtime"] = test.get("runtime", DEFAULT_RUNTIME)
+    test["exportPath"] = test.get("exportPath", DEFAULT_EXPORT_PATH)
+    test["archiveFormat"] = test.get("archiveFormat", "default")
     return test
 
 

diff --git a/model-archiver/model_archiver/tests/unit_tests/test_model_packaging_utils.py b/model-archiver/model_archiver/tests/unit_tests/test_model_packaging_utils.py
@@ -82,6 +82,7 @@ def test_archive_types(self):
             assert ar_opts.get("tgz") == ".tar.gz"
             assert ar_opts.get("no-archive") == ""
             assert ar_opts.get("default") == ".mar"
+            assert len(ar_opts) == 3
 
     # noinspection PyClassHasNoInit
     class TestCustomModelTypes: