ENH Adds default success into config plugin (#2189)

Signed-off-by: Thomas J. Fan <[email protected]>
flyteorg · Feb 14, 2024 · d840369 · d840369
1 parent 060d480
commit d840369
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git a/flytekit/clients/auth/default_html.py b/flytekit/clients/auth/default_html.py
@@ -1,4 +1,14 @@
+from contextlib import suppress
+
+
 def get_default_success_html(endpoint: str) -> str:
+    from flytekit.configuration.plugin import get_plugin
+
+    with suppress(AttributeError):
+        success_html = get_plugin().get_auth_success_html(endpoint)
+        if success_html is not None:
+            return success_html
+
     return f"""
 <html>
     <head>

diff --git a/flytekit/configuration/plugin.py b/flytekit/configuration/plugin.py
@@ -47,6 +47,10 @@ def secret_requires_group() -> bool:
     def get_default_image() -> Optional[str]:
         """Get default image. Return None to use the images from flytekit.configuration.DefaultImages"""
 
+    @staticmethod
+    def get_auth_success_html(endpoint: str) -> Optional[str]:
+        """Get default success html for auth. Return None to use flytekit's default success html."""
+
 
 class FlytekitPlugin:
     @staticmethod
@@ -80,6 +84,11 @@ def get_default_image() -> Optional[str]:
         """Get default image. Return None to use the images from flytekit.configuration.DefaultImages"""
         return None
 
+    @staticmethod
+    def get_auth_success_html(endpoint: str) -> Optional[str]:
+        """Get default success html. Return None to use flytekit's default success html."""
+        return None
+
 
 def _get_plugin_from_entrypoint():
     """Get plugin from entrypoint."""

diff --git a/tests/flytekit/unit/clients/auth/test_default_html.py b/tests/flytekit/unit/clients/auth/test_default_html.py
@@ -1,3 +1,6 @@
+from unittest.mock import Mock
+
+import flytekit
 from flytekit.clients.auth.default_html import get_default_success_html
 
 
@@ -16,3 +19,15 @@ def test_default_html():
 </html>
 """
     )  # noqa
+
+
+def test_default_html_plugin(monkeypatch):
+    def get_auth_success_html(endpoint):
+        return f"<html><head><title>Successful Auth into {endpoint}!</title></head></html>"
+
+    plugin_mock = Mock()
+    plugin_mock.get_auth_success_html.side_effect = get_auth_success_html
+    mock_global_plugin = {"plugin": plugin_mock}
+    monkeypatch.setattr(flytekit.configuration.plugin, "_GLOBAL_CONFIG", mock_global_plugin)
+
+    assert get_default_success_html("flyte.org") == get_auth_success_html("flyte.org")