[QA] Add unit tests for vizro-ai dashboard

vizro-ai/changelog.d/20240802_144558_nadija_ratkusic_graca_add_unit_tests.md

@@ -0,0 +1,48 @@
+<!--
+A new scriv changelog fragment.
+
+Uncomment the section that is right (remove the HTML comment wrapper).
+-->
+
+<!--
+### Highlights ✨
+
+- A bullet item for the Highlights ✨ category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Removed
+
+- A bullet item for the Removed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Added
+
+- A bullet item for the Added category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Changed
+
+- A bullet item for the Changed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Deprecated
+
+- A bullet item for the Deprecated category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Fixed
+
+- A bullet item for the Fixed category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->
+<!--
+### Security
+
+- A bullet item for the Security category with a link to the relevant PR at the end of your entry, e.g. Enable feature XXX ([#1](https://github.com/mckinsey/vizro/pull/1))
+
+-->

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,111 @@
+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 MockStructuredOutputLLM(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 MockStructuredOutputLLM(responses=response)
+
+
+@pytest.fixture
+def fake_llm_layout():
+    response = ['{"grid":[[0,1]]}']
+    return MockStructuredOutputLLM(responses=response)
+
+
+@pytest.fixture
+def fake_llm_filter():
+    response = ['{"column": "a", "targets": ["bar_chart"]}']
+    return MockStructuredOutputLLM(responses=response)
+
+
+@pytest.fixture
+def controllable_components():
+    return ["bar_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_cols():
+    return ["a", "b"]
+
+
+@pytest.fixture
+def df_sample(df):
+    return df.sample(5, replace=True, random_state=19)
+
+
+@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["bar_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)
+
+
+@pytest.fixture
+def filter_prompt():
+    return """
+        Create a filter from the following instructions: Filter the bar chart by column `a`.
+        Do not make up things that are optional and DO NOT configure actions, action triggers or action chains.
+        If no options are specified, leave them out."""

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,76 @@
+import pytest
+import vizro.models as vm
+from vizro.managers import model_manager
+from vizro.models import VizroBaseModel
+from vizro_ai.dashboard._response_models.controls import ControlPlan, _create_filter, _create_filter_proxy
+
+try:
+    from pydantic.v1 import ValidationError
+except ImportError:  # pragma: no cov
+    from pydantic import ValidationError
+
+# Needed for testing control creation.
+model_manager.__setitem__("bar_chart", VizroBaseModel)
+
+
+class TestFilterProxyCreate:
+    """Tests filter proxy creation."""
+
+    def test_create_filter_proxy_validate_targets(self, df_cols, df_schema, controllable_components):
+        filter_proxy = _create_filter_proxy(df_cols, df_schema, controllable_components)
+        with pytest.raises(ValidationError, match="targets must be one of"):
+            filter_proxy(targets=["population_chart"], column="a")
+
+    def test_create_filter_proxy_validate_targets_not_empty(self, df_cols, df_schema, controllable_components):
+        filter_proxy = _create_filter_proxy(df_cols=df_cols, df_schema=df_schema, controllable_components=[])
+        with pytest.raises(ValidationError):
+            filter_proxy(targets=[], column="a")
+
+    def test_create_filter_proxy_validate_columns(self, df_cols, df_schema, controllable_components):
+        filter_proxy = _create_filter_proxy(df_cols, df_schema, controllable_components)
+        with pytest.raises(ValidationError, match="column must be one of"):
+            filter_proxy(targets=["bar_chart"], column="x")
+
+    def test_create_filter_proxy(self, df_cols, df_schema, controllable_components):
+        filter_proxy = _create_filter_proxy(df_cols, df_schema, controllable_components)
+        actual_filter = filter_proxy(targets=["bar_chart"], column="a")
+
+        assert actual_filter.dict(exclude={"id": True}) == vm.Filter(targets=["bar_chart"], column="a").dict(
+            exclude={"id": True}
+        )
+
+
+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=["bar_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="bar_chart",
+            )
+
+
+def test_create_filter(filter_prompt, fake_llm_filter, df_cols, df_schema, controllable_components):
+    actual_filter = _create_filter(
+        filter_prompt=filter_prompt,
+        model=fake_llm_filter,
+        df_cols=df_cols,
+        df_schema=df_schema,
+        controllable_components=controllable_components,
+    )
+    assert actual_filter.dict(exclude={"id": True}) == vm.Filter(targets=["bar_chart"], column="a").dict(
+        exclude={"id": True}
+    )

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,9 @@
+from pandas.testing import assert_frame_equal
+from vizro_ai.dashboard._response_models.df_info import _get_df_info
+
+
+def test_get_df_info(df, df_schema, df_sample):
+    actual_df_schema, actual_df_sample = _get_df_info(df=df)
+
+    assert actual_df_schema == df_schema
+    assert_frame_equal(actual_df_sample, df_sample)