[Test] Update vizro ai integration test for dataclass output

vizro-ai/changelog.d/20240606_155107_anna_xiong_update_vizro_ai_integration_test.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/src/vizro_ai/_vizro_ai.py

@@ -1,5 +1,4 @@
 import logging
-from dataclasses import asdict
 from typing import Any, Optional, Union
 
 import pandas as pd
@@ -25,7 +24,6 @@ class VizroAI:
     """Vizro-AI main class."""
 
     pipeline_manager: PipelineManager = PipelineManager()
-    _return_all_text: bool = False  # TODO deleted after adding new integration test
 
     def __init__(self, model: Optional[Union[ChatOpenAI, str]] = None):
         """Initialization of VizroAI.
@@ -153,10 +151,4 @@ def plot(  # pylint: disable=too-many-arguments  # noqa: PLR0913
                 code_explain=vizro_plot.code_explanation,
             )
 
-        # TODO Tentative for integration test, will be updated/removed for new tests
-        if self._return_all_text:
-            output_dict = asdict(vizro_plot)
-            output_dict["code_string"] = vizro_plot.code
-            return output_dict
-
         return vizro_plot if return_elements else vizro_plot.figure

vizro-ai/tests/integration/test_example.py

@@ -7,44 +7,48 @@
 
 
 def test_chart():
-    vizro_ai._return_all_text = True
-    resp = vizro_ai.plot(df, "describe the composition of scatter chart with gdp in continent")
+    resp = vizro_ai.plot(
+        df=df,
+        user_input="describe the composition of scatter chart with gdp in continent",
+        explain=False,
+        return_elements=True,
+    )
     assert_that(
-        resp["code_string"],
+        resp.code,
         all_of(contains_string("px.scatter")),
     )
     assert_that(
-        resp["code_string"],
+        resp.code,
         any_of(contains_string("x='continent'"), contains_string("x='gdpPercap'")),
     )
     assert_that(
-        resp["code_string"],
+        resp.code,
         any_of(contains_string("y='count'"), contains_string("y='gdpPercap'"), contains_string("y='continent'")),
     )
-    assert_that(resp["business_insights"], equal_to(None))
-    assert_that(resp["code_explanation"], equal_to(None))
+    assert_that(resp.code_explanation, equal_to(None))
+    assert_that(resp.business_insights, equal_to(None))
 
 
 def test_chart_with_explanation():
     vizro_ai._return_all_text = True
-    resp = vizro_ai.plot(df, "describe the composition of gdp in US", explain=True)
+    resp = vizro_ai.plot(df, "describe the composition of gdp in US", explain=True, return_elements=True)
     assert_that(
-        resp["code_string"],
+        resp.code,
         all_of(contains_string("px.bar"), contains_string("x='year'")),
     )
     assert_that(
-        resp["code_string"],
+        resp.code,
         any_of(contains_string("y='gdpPercap'"), contains_string("y='total_gdp'")),
     )
     assert_that(
-        resp["business_insights"],
+        resp.business_insights,
         any_of(
             contains_string("GDP per capita in the United States"),
             contains_string("GDP in the United States"),
             contains_string("GDP in the US"),
         ),
     )
     assert_that(
-        resp["code_explanation"],
+        resp.code_explanation,
         all_of(contains_string("https://vizro.readthedocs.io/en/stable/pages/user_guides/custom_charts/")),
     )