more tests

Signed-off-by: Kevin Su <[email protected]>

flytekit/types/structured/structured_dataset.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import collections
-import importlib
 import os
 import types
 import typing

plugins/flytekit-polars/tests/test_polars_plugin_sd.py

@@ -1,4 +1,6 @@
+import pandas as pd
 import polars as pl
+from flytekitplugins.polars.sd_transformers import PolarsDataFrameRenderer
 from typing_extensions import Annotated
 
 from flytekit import kwtypes, task, workflow
@@ -57,3 +59,10 @@ def wf() -> full_schema:
 
     result = wf()
     assert result is not None
+
+
+def test_polars_renderer():
+    df = pl.DataFrame({"col1": [1, 3, 2], "col2": list("abc")})
+    assert PolarsDataFrameRenderer().to_html(df) == pd.DataFrame(
+        df.describe().transpose(), columns=df.describe().columns
+    ).to_html(index=False)

tests/flytekit/unit/core/test_structured_dataset.py

@@ -414,3 +414,18 @@ def test_protocol_detection():
 
         protocol = e._protocol_from_type_or_prefix(ctx2, pd.DataFrame, "bq://foo")
         assert protocol == "bq"
+
+
+def test_register_renderers():
+    class DummyRenderer:
+        def to_html(self, input: str) -> str:
+            return "hello " + input
+
+    renderers = StructuredDatasetTransformerEngine.Renderers
+    StructuredDatasetTransformerEngine.register_renderer(str, DummyRenderer())
+    assert renderers[str].to_html("flyte") == "hello flyte"
+    assert pd.DataFrame in renderers
+    assert pa.Table in renderers
+
+    with pytest.raises(NotImplementedError, match="Could not find a renderer for <class 'int'> in"):
+        StructuredDatasetTransformerEngine().to_html(FlyteContextManager.current_context(), 3, int)

tests/flytekit/unit/deck/test_renderer.py

@@ -1,9 +1,12 @@
 import pandas as pd
+import pyarrow as pa
 
-from flytekit.deck.renderer import TopFrameRenderer
+from flytekit.deck.renderer import ArrowRenderer, TopFrameRenderer
 
 
-def test_frame_profiling_renderer():
+def test_renderer():
     df = pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [1, 22]})
-    renderer = TopFrameRenderer()
-    assert renderer.to_html(df) == df.to_html()
+    pa_df = pa.Table.from_pandas(df)
+
+    assert TopFrameRenderer().to_html(df) == df.to_html()
+    assert ArrowRenderer().to_html(pa_df) == pa_df.to_string()