Add Flyte Backend Version to pyflyte info command (#2938)

* Add backend version line to pyflyte info command Signed-off-by: davidlin20dev <[email protected]> * test: Add pyflyte info command versoin tests Signed-off-by: davidlin20dev <[email protected]> --------- Signed-off-by: davidlin20dev <[email protected]>
flyteorg · Nov 19, 2024 · 39e7227 · 39e7227
1 parent ee4f9ab
commit 39e7227
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 4 deletions.
diff --git a/flytekit/clis/version.py b/flytekit/clis/version.py
@@ -8,7 +8,8 @@
 Content = """
 This CLI is meant to be used within a virtual environment that has Flytekit installed. Ideally it is used to iterate on your Flyte workflows and tasks.
 
-Flytekit Version: [cyan]{version}[reset]
+Flytekit Version: [cyan]{flytekit_version}[reset]
+Flyte Backend Version: [cyan]{backend_version}[reset]
 Flyte Backend Endpoint: [cyan]{endpoint}
 """
 
@@ -17,11 +18,14 @@
 @click.pass_context
 def info(ctx: click.Context):
     """
-    Print out information about the current Flyte Python CLI environment - like the version of Flytekit, backend endpoint
-    currently configured, etc.
+    Print out information about the current Flyte Python CLI environment - like the version of Flytekit, the version of Flyte Backend Version,
+    backend endpoint currently configured, etc.
     """
     import flytekit
 
     remote: FlyteRemote = get_and_save_remote_with_click_context(ctx, project="flytesnacks", domain="development")
-    c = Content.format(version=flytekit.__version__, endpoint=remote.client.url)
+    backend_version = remote.client.get_control_plane_version()
+    c = Content.format(
+        flytekit_version=flytekit.__version__, backend_version=backend_version, endpoint=remote.client.url
+    )
     rich.print(Panel(c, title="Flytekit CLI Info", border_style="purple", padding=(1, 1, 1, 1)))
diff --git a/tests/flytekit/integration/clis/test_version.py b/tests/flytekit/integration/clis/test_version.py
@@ -0,0 +1,32 @@
+import subprocess
+import re
+
+def run_info_command() -> str:
+    """
+    Runs the `pyflyte info` command and returns its output.
+    """
+    out = subprocess.run(
+        ["pyflyte", "info"],
+        capture_output=True,  # Capture the output streams
+        text=True,  # Return outputs as strings (not bytes)
+    )
+    # Ensure the command ran successfully
+    assert out.returncode == 0, (f"Command failed with return code {out.returncode}.\n"
+                                 f"Standard Output: {out.stdout}\n"
+                                 f"Standard Error: {out.stderr}\n")
+    return out.stdout
+
+def test_info_command_versions():
+    """
+    Tests that the `pyflyte info` command outputs the correct version information.
+    """
+    output = run_info_command()
+
+    # Check that Flytekit Version is displayed
+    assert re.search(r"Flytekit Version: \S+", output), "Flytekit Version not found in output."
+
+    # Check that Flyte Backend Version is displayed
+    assert re.search(r"Flyte Backend Version: \S+", output), "Flyte Backend Version not found in output."
+
+    # Check that Flyte Backend Endpoint is displayed
+    assert re.search(r"Flyte Backend Endpoint: \S+", output), "Flyte Backend Endpoint not found in output."