adding unit tests for vizro-ai dashboard

vizro-ai/tests/unit/vizro-ai/dashboard/_graph/conftest.py

@@ -0,0 +1,31 @@
+import pandas as pd
+import pytest
+from langchain_core.messages import HumanMessage
+from vizro_ai.dashboard._graph.dashboard_creation import GraphState
+from vizro_ai.dashboard.utils import AllDfMetadata, DfMetadata
+
+
+@pytest.fixture
+def dataframes():
+    return [pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]})]
+
+
+@pytest.fixture
+def df_metadata():
+    df_metadata = AllDfMetadata({})
+    df_metadata.all_df_metadata["gdp_chart"] = DfMetadata(
+        df_schema={"a": "int64", "b": "int64"},
+        df=pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]}),
+        df_sample=pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]}),
+    )
+    return df_metadata
+
+
+@pytest.fixture
+def graph_state(dataframes, df_metadata):
+    return GraphState(
+        messages=[HumanMessage(content="contents of the message")],
+        dfs=dataframes,
+        all_df_metadata=df_metadata,
+        pages=[],
+    )

vizro-ai/tests/unit/vizro-ai/dashboard/_graph/test_dashboard_creation.py

@@ -0,0 +1,37 @@
+import pandas as pd
+import pytest
+
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:  # pragma: no cov
+    from pydantic import ValidationError
+
+from langchain_core.messages import HumanMessage
+from vizro_ai.dashboard._graph.dashboard_creation import GraphState
+
+
+class TestConfig:
+    """Test GraphState config creation."""
+
+    def test_graph_state_instantiation(self, graph_state, dataframes):
+        assert isinstance(graph_state, GraphState)
+        assert graph_state.messages[0].content == "contents of the message"
+        assert graph_state.dfs == dataframes
+        assert "gdp_chart" in graph_state.all_df_metadata.all_df_metadata
+        assert graph_state.pages == []
+
+    @pytest.mark.parametrize(
+        "dataframes, output_error",
+        [
+            (pd.DataFrame(), "value is not a valid list"),
+            ([pd.DataFrame(), {}], "instance of DataFrame expected"),
+        ],
+    )
+    def test_check_dataframes(self, dataframes, output_error, df_metadata):
+        with pytest.raises(ValidationError, match=output_error):
+            GraphState(
+                messages=[HumanMessage(content="contents of the message")],
+                dfs=dataframes,
+                all_df_metadata=df_metadata,
+                pages=[],
+            )

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/conftest.py

@@ -0,0 +1,103 @@
+from typing import Any, List
+
+import pandas as pd
+import pytest
+from langchain.output_parsers import PydanticOutputParser
+from langchain_community.llms.fake import FakeListLLM
+from vizro_ai.dashboard._response_models.components import ComponentPlan
+from vizro_ai.dashboard._response_models.page import PagePlan
+from vizro_ai.dashboard.utils import AllDfMetadata, DfMetadata
+
+
+class FakeListLLM(FakeListLLM):
+    def bind_tools(self, tools: List[Any]):
+        return super().bind(tools=tools)
+
+    def with_structured_output(self, schema):
+        llm = self
+        output_parser = PydanticOutputParser(pydantic_object=schema)
+        return llm | output_parser
+
+
+@pytest.fixture
+def fake_llm_card():
+    response = ['{"text":"this is a card","href":""}']
+    return FakeListLLM(responses=response)
+
+
+@pytest.fixture
+def fake_llm_layout():
+    response = ['{"grid":[[0,1]]}']
+    return FakeListLLM(responses=response)
+
+
+@pytest.fixture
+def fake_llm_filter():
+    response = ['{"column": "a", "targets": ["gdp_chart"]}']
+    return FakeListLLM(responses=response)
+
+
+@pytest.fixture
+def df_cols():
+    return ["continent", "country", "population", "gdp"]
+
+
+@pytest.fixture
+def controllable_components():
+    return ["gdp_chart"]
+
+
+@pytest.fixture
+def layout_description():
+    return "The layout of this page should use `grid=[[0,1]]`"
+
+
+@pytest.fixture
+def df():
+    return pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]})
+
+
+@pytest.fixture
+def df_sample():
+    return pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]})
+
+
+@pytest.fixture
+def df_schema():
+    return {"a": "int64", "b": "int64"}
+
+
+@pytest.fixture
+def df_metadata(df, df_schema, df_sample):
+    df_metadata = AllDfMetadata({})
+    df_metadata.all_df_metadata["gdp_chart"] = DfMetadata(
+        df_schema=df_schema,
+        df=df,
+        df_sample=df_sample,
+    )
+    return df_metadata
+
+
+@pytest.fixture
+def component_card():
+    return ComponentPlan(
+        component_type="Card",
+        component_description="This is a card",
+        component_id="card_1",
+        df_name="N/A",
+    )
+
+
+@pytest.fixture
+def component_card_2():
+    return ComponentPlan(
+        component_type="Card",
+        component_description="This is a second card",
+        component_id="card_2",
+        df_name="N/A",
+    )
+
+
+@pytest.fixture
+def page_plan(component_card):
+    return PagePlan(title="Test Page", components_plan=[component_card], controls_plan=[], layout_plan=None)

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_components.py

@@ -0,0 +1,27 @@
+import pytest
+from vizro_ai.dashboard._response_models.components import ComponentPlan
+
+
+class TestComponentCreate:
+    """Tests component creation."""
+
+    def test_component_plan_instantiation(self):
+        component = ComponentPlan(
+            component_id="card_1",
+            component_type="Card",
+            component_description="This is a card",
+            df_name="N/A",
+        )
+        assert component.component_id == "card_1"
+        assert component.component_type == "Card"
+        assert component.component_description == "This is a card"
+        assert component.df_name == "N/A"
+
+    @pytest.mark.xfail(raises=ValueError, reason="Known issue: real model is required for .plot")
+    def test_card_create(self, component_card, fake_llm_card):
+        if component_card.component_type == "Card":
+            actual = component_card.create(
+                model=fake_llm_card,
+                all_df_metadata=None,
+            )
+        assert actual.type == "card"

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_controls.py

@@ -0,0 +1,49 @@
+import pytest
+from vizro_ai.dashboard._response_models.controls import ControlPlan, _create_filter_proxy
+
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:  # pragma: no cov
+    from pydantic import ValidationError
+
+
+class TestControlCreate:
+    """Tests control creation."""
+
+    def test_create_filter_proxy_validate_targets(self, df_cols, df_schema, controllable_components):
+        actual = _create_filter_proxy(df_cols, df_schema, controllable_components)
+        with pytest.raises(ValidationError, match="targets must be one of"):
+            actual(targets=["population_chart"], column="gdp")
+
+    def test_create_filter_proxy_validate_targets_not_empty(self, df_cols, df_schema, controllable_components):
+        actual = _create_filter_proxy(df_cols=df_cols, df_schema=df_schema, controllable_components=[])
+        with pytest.raises(ValidationError):
+            actual(targets=[], column="gdp")
+
+    def test_create_filter_proxy_validate_columns(self, df_cols, df_schema, controllable_components):
+        actual = _create_filter_proxy(df_cols, df_schema, controllable_components)
+        with pytest.raises(ValidationError, match="column must be one of"):
+            actual(targets=["gdp_chart"], column="x")
+
+
+class TestControlPlan:
+    """Test control plan."""
+
+    def test_control_plan_invalid_df_name(self, fake_llm_filter, df_metadata):
+        control_plan = ControlPlan(
+            control_type="Filter",
+            control_description="Create a filter that filters the data based on the column 'a'.",
+            df_name="population_chart",
+        )
+        default_control = control_plan.create(
+            model=fake_llm_filter, controllable_components=["gdp_chart"], all_df_metadata=df_metadata
+        )
+        assert default_control is None
+
+    def test_control_plan_invalid_type(self, fake_llm_filter, df_metadata):
+        with pytest.raises(ValidationError):
+            ControlPlan(
+                control_type="parameter",
+                control_description="Create a parameter that targets the data based on the column 'a'.",
+                df_name="gdp_chart",
+            )

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_dashboard.py

@@ -0,0 +1,18 @@
+from vizro_ai.dashboard._response_models.dashboard import DashboardPlan
+
+
+class TestDashboardPlanner:
+    """Tests dashboard planner."""
+
+    def test_dashboard_planner(self, page_plan):
+        dashboard_plan = DashboardPlan(
+            title="Test Dashboard",
+            pages=[page_plan],
+        )
+        assert dashboard_plan.pages[0].title == "Test Page"
+        assert dashboard_plan.pages[0].components_plan[0].component_id == "card_1"
+        assert dashboard_plan.pages[0].components_plan[0].component_type == "Card"
+        assert dashboard_plan.pages[0].components_plan[0].component_description == "This is a card"
+        assert dashboard_plan.pages[0].components_plan[0].df_name == "N/A"
+        assert dashboard_plan.pages[0].layout_plan is None
+        assert dashboard_plan.pages[0].controls_plan == []

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_df_info.py

@@ -0,0 +1,7 @@
+from vizro_ai.dashboard._response_models.df_info import _get_df_info
+
+
+def test_get_df_info(df, df_schema):
+    actual_df_schema, _ = _get_df_info(df=df)
+
+    assert actual_df_schema == df_schema

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_layout.py

@@ -0,0 +1,43 @@
+import pytest
+import vizro.models as vm
+from vizro_ai.dashboard._pydantic_output import _get_pydantic_model
+from vizro_ai.dashboard._response_models.layout import LayoutPlan, _convert_to_grid
+
+
+class TestLayoutPlan:
+    """Test layout creation."""
+
+    def test_structured_output_layout_create(self, fake_llm_layout, layout_description):
+        structured_output = _get_pydantic_model(
+            query=layout_description, llm_model=fake_llm_layout, response_model=vm.Layout, df_info=None
+        )
+        assert structured_output.dict(exclude={"id": True}) == vm.Layout(grid=[[0, 1]]).dict(exclude={"id": True})
+
+    def test_layout_plan(self):
+        layout_plan = LayoutPlan(
+            layout_grid_template_areas=["graph card"],
+        )
+        layout = layout_plan.create(["graph", "card"])
+
+        assert layout.dict(exclude={"id": True}) == vm.Layout(grid=[[0, 1]]).dict(exclude={"id": True})
+
+
+@pytest.mark.parametrize(
+    "layout_grid_template_areas, component_ids, grid",
+    [
+        (
+            ["card_1 scatter_plot scatter_plot", "card_2 scatter_plot scatter_plot"],
+            ["card_1", "scatter_plot", "card_2"],
+            [[0, 1, 1], [2, 1, 1]],
+        ),
+        (
+            ["card_1 scatter_plot scatter_plot", "card_2 scatter_plot scatter_plot"],
+            ["card_1", "scatter_plot"],
+            [],
+        ),
+    ],
+)
+def test_convert_to_grid(layout_grid_template_areas, component_ids, grid):
+    actual_grid = _convert_to_grid(layout_grid_template_areas=layout_grid_template_areas, component_ids=component_ids)
+
+    assert actual_grid == grid

vizro-ai/tests/unit/vizro-ai/dashboard/_response_models/test_page.py

@@ -0,0 +1,55 @@
+import pytest
+from vizro_ai.dashboard._response_models.page import PagePlan
+
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:  # pragma: no cov
+    from pydantic import ValidationError
+
+
+class TestPagePlan:
+    """Test for page plan."""
+
+    def test_dashboard_plan(self, component_card):
+        page_plan = PagePlan(
+            title="Test Page",
+            components_plan=[component_card],
+            controls_plan=[],
+            layout_plan=None,
+        )
+        assert page_plan.title == "Test Page"
+        assert page_plan.components_plan[0].component_id == "card_1"
+        assert page_plan.components_plan[0].component_type == "Card"
+        assert page_plan.components_plan[0].component_description == "This is a card"
+        assert page_plan.layout_plan is None
+        assert page_plan.controls_plan == []
+        assert page_plan.unsupported_specs == []
+
+    def test_page_plan_invalid_components(self):
+        with pytest.raises(ValidationError, match="A page must contain at least one component."):
+            PagePlan(
+                title="Test Page",
+                components_plan=[],
+                controls_plan=[],
+                layout_plan=None,
+            )
+
+    def test_page_plan_unsupported_specs(self, component_card):
+        page_plan = PagePlan(
+            title="Test Page",
+            components_plan=[component_card],
+            controls_plan=[],
+            layout_plan=None,
+            unsupported_specs=["Unknown"],
+        )
+
+        assert page_plan.unsupported_specs == []
+
+    def test_page_plan_duplicate_components(self, component_card):
+        with pytest.raises(ValidationError):
+            PagePlan(
+                title="Test Page",
+                components_plan=[component_card, component_card],
+                controls_plan=[],
+                layout_plan=None,
+            )