From db9f744c37b1fe06abdb22c3682234cc32c024cf Mon Sep 17 00:00:00 2001
From: paullongtan <paullongtan@gmail.com>
Date: Wed, 1 Jan 2025 22:08:29 -0800
Subject: [PATCH] test: add yaml output format and registration without
 specifying --summary-dir testings to unit test

Signed-off-by: paullongtan <paullongtan@gmail.com>
---
 .../unit/cli/pyflyte/test_register.py         | 85 +++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/tests/flytekit/unit/cli/pyflyte/test_register.py b/tests/flytekit/unit/cli/pyflyte/test_register.py
index 94c7028935..f1e627cfcc 100644
--- a/tests/flytekit/unit/cli/pyflyte/test_register.py
+++ b/tests/flytekit/unit/cli/pyflyte/test_register.py
@@ -2,6 +2,7 @@
 import shutil
 import subprocess
 import json
+import yaml
 
 import mock
 import pytest
@@ -218,6 +219,27 @@ def test_register_registrated_summary_json(mock_client, mock_remote):
         summary_path = os.path.join("summaries", "registration_summary.json")
         assert os.path.exists(summary_path)
 
+        with open(summary_path) as f:
+            summary_data = json.load(f)
+
+        assert isinstance(summary_data, list)
+        assert len(summary_data) > 0
+        for entry in summary_data:
+            assert "id" in entry
+            assert "type" in entry
+            assert "version" in entry
+            assert "status" in entry
+
+
+        result = runner.invoke(
+            pyflyte.main,
+            ["register", "--summary-format", "json", "core5"]
+        )
+        assert result.exit_code == 0
+
+        summary_path = os.path.join(os.getcwd(), "registration_summary.json")
+        assert os.path.exists(summary_path)
+
         with open(summary_path) as f:
             summary_data = json.load(f)
 
@@ -231,3 +253,66 @@ def test_register_registrated_summary_json(mock_client, mock_remote):
 
         shutil.rmtree("core5")
         shutil.rmtree("summaries")
+
+@mock.patch("flytekit.configuration.plugin.FlyteRemote", spec=FlyteRemote)
+@mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
+def test_register_registrated_summary_yaml(mock_client, mock_remote):
+    ctx = FlyteContextManager.current_context()
+    mock_remote._client = mock_client
+    mock_remote.return_value.context = ctx
+    mock_remote.return_value._version_from_hash.return_value = "dummy_version_from_hash"
+    mock_remote.return_value.fast_package.return_value = "dummy_md5_bytes", "dummy_native_url"
+    runner = CliRunner()
+    context_manager.FlyteEntities.entities.clear()
+
+    with runner.isolated_filesystem():
+        out = subprocess.run(["git", "init"], capture_output=True)
+        assert out.returncode == 0
+        os.makedirs("core6", exist_ok=True)
+        os.makedirs("summaries", exist_ok=True)
+        with open(os.path.join("core6", "sample.py"), "w") as f:
+            f.write(sample_file_contents)
+            f.close()
+
+        result = runner.invoke(
+            pyflyte.main,
+            ["register", "--summary-format", "yaml", "--summary-dir", "summaries", "core6"]
+        )
+        assert result.exit_code == 0
+
+        summary_path = os.path.join("summaries", "registration_summary.yaml")
+        assert os.path.exists(summary_path)
+
+        with open(summary_path) as f:
+            summary_data = yaml.safe_load(f)
+
+        assert isinstance(summary_data, list)
+        assert len(summary_data) > 0
+        for entry in summary_data:
+            assert "id" in entry
+            assert "type" in entry
+            assert "version" in entry
+            assert "status" in entry
+
+        result = runner.invoke(
+            pyflyte.main,
+            ["register", "--summary-format", "yaml", "core6"]
+        )
+        assert result.exit_code == 0
+
+        summary_path = os.path.join(os.getcwd(), "registration_summary.yaml")
+        assert os.path.exists(summary_path)
+
+        with open(summary_path) as f:
+            summary_data = yaml.safe_load(f)
+
+        assert isinstance(summary_data, list)
+        assert len(summary_data) > 0
+        for entry in summary_data:
+            assert "id" in entry
+            assert "type" in entry
+            assert "version" in entry
+            assert "status" in entry
+
+        shutil.rmtree("core6")
+        shutil.rmtree("summaries")