fix: fix unit test

flytekit/tools/repo.py

@@ -399,11 +399,13 @@ async def _register(entities: typing.List[task.TaskSpec]):
     click.secho(f"Successfully registered {len(registrable_entities)} entities", fg="green")
 
     if summary_format is not None:
+        supported_foramt = ["json", "yaml"]
+        if summary_format not in supported_foramt:
+            raise ValueError(f"Unsupported file format: {summary_format}")
+
         if summary_dir is not None:
-            # Directory path is already absolute and resolved via click.Path
             os.makedirs(summary_dir, exist_ok=True)
         else:
-            # Default to current working directory if not specified
             summary_dir = os.getcwd()
 
         summary_file = f"registration_summary.{summary_format}"
@@ -415,7 +417,5 @@ async def _register(entities: typing.List[task.TaskSpec]):
         elif summary_format == "yaml":
             with open(summary_path, "w") as f:
                 yaml.dump(all_results, f)
-        else:
-            raise ValueError(f"Unsupported file format: {summary_format}")
 
         click.secho(f"Registration summary written to: {summary_path}", fg="green")

haha/registration_summary.json

@@ -0,0 +1 @@
+[{"id": "test.sum", "type": "TASK", "version": "1C7xMt9aD2YVbcFO6N0hgw", "status": "success"}, {"id": "test.square", "type": "TASK", "version": "1C7xMt9aD2YVbcFO6N0hgw", "status": "success"}, {"id": "test.my_workflow", "type": "WORKFLOW", "version": "1C7xMt9aD2YVbcFO6N0hgw", "status": "success"}, {"id": "test.my_workflow", "type": "LAUNCH_PLAN", "version": "1C7xMt9aD2YVbcFO6N0hgw", "status": "success"}]

registration_summary.json

@@ -0,0 +1 @@
+[{"id": "tmp_test.test.sum", "type": "TASK", "version": "-uRk_xLMhDS-jsXtewBKcA", "status": "success"}, {"id": "tmp_test.test.square", "type": "TASK", "version": "-uRk_xLMhDS-jsXtewBKcA", "status": "success"}, {"id": "tmp_test.test.my_workflow", "type": "WORKFLOW", "version": "-uRk_xLMhDS-jsXtewBKcA", "status": "success"}, {"id": "tmp_test.test.my_workflow", "type": "LAUNCH_PLAN", "version": "-uRk_xLMhDS-jsXtewBKcA", "status": "success"}]

tests/flytekit/unit/cli/pyflyte/test_register.py

@@ -194,10 +194,11 @@ def test_register_summary_dir_without_format(mock_client, mock_remote):
 @mock.patch("flytekit.configuration.plugin.FlyteRemote", spec=FlyteRemote)
 @mock.patch("flytekit.clients.friendly.SynchronousFlyteClient", spec=SynchronousFlyteClient)
 def test_register_registrated_summary_json(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()
 
@@ -211,15 +212,14 @@ def test_register_registrated_summary_json(mock_client, mock_remote):
             f.close()
 
         # Run registration command
-        # result = runner.invoke(
-        #     pyflyte.main,
-        #     ["register", "--summary-format", "json", "--summary-dir", "summaries", "core5"]
-        # )
         result = runner.invoke(
             pyflyte.main,
-            ["register", "--summary-format", "json", "core5"]
+            ["register", "--summary-format", "json", "--summary-dir", "summaries", "core5"]
         )
-
+        # result = runner.invoke(
+        #     pyflyte.main,
+        #     ["register", "--summary-format", "json", "core5"]
+        # )
         assert result.exit_code == 0
 
         summary_path = os.path.join("summaries", "registration_summary.json")

tmp_test/registration_summary.json

@@ -0,0 +1 @@
+[{"id": "test.sum", "type": "TASK", "version": "hJdmkFl7o-uuiXap9L7_fg", "status": "success"}, {"id": "test.square", "type": "TASK", "version": "hJdmkFl7o-uuiXap9L7_fg", "status": "success"}, {"id": "test.my_workflow", "type": "WORKFLOW", "version": "hJdmkFl7o-uuiXap9L7_fg", "status": "success"}, {"id": "test.my_workflow", "type": "LAUNCH_PLAN", "version": "hJdmkFl7o-uuiXap9L7_fg", "status": "success"}]

tmp_test/registration_summary.yaml

@@ -0,0 +1,16 @@
+- id: test.sum
+  status: success
+  type: TASK
+  version: 1C7xMt9aD2YVbcFO6N0hgw
+- id: test.square
+  status: success
+  type: TASK
+  version: 1C7xMt9aD2YVbcFO6N0hgw
+- id: test.my_workflow
+  status: success
+  type: WORKFLOW
+  version: 1C7xMt9aD2YVbcFO6N0hgw
+- id: test.my_workflow
+  status: success
+  type: LAUNCH_PLAN
+  version: 1C7xMt9aD2YVbcFO6N0hgw

tmp_test/test.py

@@ -0,0 +1,13 @@
+from flytekit import task, workflow
+
+@task
+def sum(x: int, y: int) -> int:
+    return x + y
+
+@task
+def square(z: int) -> int:
+    return z*z
+
+@workflow
+def my_workflow(x: int, y: int) -> int:
+    return sum(x=square(z=x), y=square(z=y))