Introduce a command group that allows adding a command to generate se…

…rvice account id tokens later

Signed-off-by: Fabio Grätz <[email protected]>

plugins/flytekit-identity-aware-proxy/flytekitplugins/identity_aware_proxy/cli.py

@@ -99,7 +99,13 @@ def get_gcp_secret_manager_secret(project_id: str, secret_id: str, version: typi
     return payload
 
 
-@click.command()
+@click.group()
+def cli():
+    """Generate ID tokens for GCP Identity Aware Proxy (IAP)."""
+    pass
+
+
+@cli.command()
 @click.option(
     "--desktop_client_id",
     type=str,
@@ -138,10 +144,10 @@ def get_gcp_secret_manager_secret(project_id: str, secret_id: str, version: typi
     required=True,
     help="GCP project ID (in which `desktop_client_secret_gcp_secret_name` is saved).",
 )
-def flyte_iap_token(
+def generate_user_id_token(
     desktop_client_id: str, desktop_client_secret_gcp_secret_name: str, webapp_client_id: str, project: str
 ):
-    """Generate an ID token for proxy-authentication/authorization with GCP Identity Aware Proxy."""
+    """Generate a user account ID token for proxy-authentication/authorization with GCP Identity Aware Proxy."""
     desktop_client_secret = get_gcp_secret_manager_secret(project, desktop_client_secret_gcp_secret_name)
 
     iap_authenticator = GCPIdentityAwareProxyAuthenticator(
@@ -158,4 +164,4 @@ def flyte_iap_token(
 
 
 if __name__ == "__main__":
-    flyte_iap_token()
+    cli()

plugins/flytekit-identity-aware-proxy/setup.py

@@ -21,7 +21,7 @@
     packages=[f"flytekitplugins.{PLUGIN_NAME}"],
     entry_points={
         "console_scripts": [
-            "flyte-iap-token=flytekitplugins.identity_aware_proxy.cli:flyte_iap_token",
+            "flyte-iap=flytekitplugins.identity_aware_proxy.cli:cli",
         ],
     },
     install_requires=plugin_requires,

plugins/flytekit-identity-aware-proxy/tests/test_flytekitplugins_iap.py

@@ -3,15 +3,19 @@
 import click
 import pytest
 from click.testing import CliRunner
-from flytekitplugins.identity_aware_proxy.cli import flyte_iap_token, get_gcp_secret_manager_secret
+from flytekitplugins.identity_aware_proxy.cli import cli, get_gcp_secret_manager_secret
 from google.api_core.exceptions import NotFound
 
 
 def test_help() -> None:
     """Smoke test external command IAP ID token generator cli by printing help message."""
     runner = CliRunner()
-    result = runner.invoke(flyte_iap_token, "--help")
-    assert "Generate an ID token" in result.output
+    result = runner.invoke(cli, "--help")
+    assert "Generate ID tokens" in result.output
+    assert result.exit_code == 0
+
+    result = runner.invoke(cli, ["generate-user-id-token", "--help"])
+    assert "Generate a user account ID token" in result.output
     assert result.exit_code == 0