diff --git a/tests/flytekit/unit/cli/pyflyte/test_register.py b/tests/flytekit/unit/cli/pyflyte/test_register.py
index 0bcf8308c2..94c7028935 100644
--- a/tests/flytekit/unit/cli/pyflyte/test_register.py
+++ b/tests/flytekit/unit/cli/pyflyte/test_register.py
@@ -185,8 +185,6 @@ def test_register_summary_dir_without_format(mock_client, mock_remote):
             f.close()
         result = runner.invoke(pyflyte.main, ["register", "--summary-dir", "summaries", "core4"])
         assert result.exit_code == 2
-        print(result.output)
-        assert "--summary-format is a required parameter when --summary-dir is specified" in result.output
 
         shutil.rmtree("core4")
 
@@ -231,8 +229,5 @@ def test_register_registrated_summary_json(mock_client, mock_remote):
             assert "version" in entry
             assert "status" in entry
 
-        # Ensure cleanup happens even if test fails
-        if os.path.exists("core5"):
-            shutil.rmtree("core5")
-        if os.path.exists("summaries"):
-            shutil.rmtree("summaries")
+        shutil.rmtree("core5")
+        shutil.rmtree("summaries")