From 7de270a89ba33823b108ef71e0362f311540bb78 Mon Sep 17 00:00:00 2001
From: Alexey Snigir <alexey_snigir@external.mckinsey.com>
Date: Mon, 5 Aug 2024 13:49:12 +0200
Subject: [PATCH] all tests run

---
 .../workflows/test-integration-vizro-ai.yml   |  20 ++-
 vizro-ai/tests/integration/test_example.py    | 120 +++++++++---------
 2 files changed, 69 insertions(+), 71 deletions(-)

diff --git a/.github/workflows/test-integration-vizro-ai.yml b/.github/workflows/test-integration-vizro-ai.yml
index da291114d..2dbaed8b3 100644
--- a/.github/workflows/test-integration-vizro-ai.yml
+++ b/.github/workflows/test-integration-vizro-ai.yml
@@ -80,18 +80,16 @@ jobs:
 
       - name: Run vizro-ai integration tests with pypi vizro
         run: |
-          export LANGCHAIN_API_KEY=${{ secrets.LANGCHAIN_API_KEY }}
-          export LANGCHAIN_TRACING_V2=false
           export OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}
           export OPENAI_API_BASE=${{ secrets.OPENAI_API_BASE }}
           hatch run ${{ matrix.hatch-env }}:test-integration
 
-#      - name: Run vizro-ai integration tests with local vizro
-#        run: |
-#          export OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}
-#          export OPENAI_API_BASE=${{ secrets.OPENAI_API_BASE }}
-#          cd ../vizro-core
-#          hatch build
-#          cd ../vizro-ai
-#          hatch run ${{ matrix.hatch-env }}:pip install ../vizro-core/dist/vizro*.tar.gz
-#          hatch run ${{ matrix.hatch-env }}:test-integration
+      - name: Run vizro-ai integration tests with local vizro
+        run: |
+          export OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }}
+          export OPENAI_API_BASE=${{ secrets.OPENAI_API_BASE }}
+          cd ../vizro-core
+          hatch build
+          cd ../vizro-ai
+          hatch run ${{ matrix.hatch-env }}:pip install ../vizro-core/dist/vizro*.tar.gz
+          hatch run ${{ matrix.hatch-env }}:test-integration
diff --git a/vizro-ai/tests/integration/test_example.py b/vizro-ai/tests/integration/test_example.py
index 5ff4368d5..9c9ad64d6 100644
--- a/vizro-ai/tests/integration/test_example.py
+++ b/vizro-ai/tests/integration/test_example.py
@@ -1,60 +1,60 @@
-# import vizro.plotly.express as px
-# from hamcrest import all_of, any_of, assert_that, contains_string, equal_to
-# from vizro_ai import VizroAI
-#
-# vizro_ai = VizroAI()
-# df = px.data.gapminder()
-#
-#
-# def test_chart():
-#     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,
-#         all_of(contains_string("px.scatter")),
-#     )
-#     assert_that(
-#         resp.code,
-#         any_of(contains_string("x='continent'"), contains_string("x='gdpPercap'")),
-#     )
-#     assert_that(
-#         resp.code,
-#         any_of(contains_string("y='count'"), contains_string("y='gdpPercap'"), contains_string("y='continent'")),
-#     )
-#     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, return_elements=True)
-#     assert_that(
-#         resp.code,
-#         all_of(contains_string("px.bar"), contains_string("x='year'")),
-#     )
-#     assert_that(
-#         resp.code,
-#         any_of(contains_string("y='gdpPercap'"), contains_string("y='total_gdp'")),
-#     )
-#     assert_that(
-#         resp.business_insights,
-#         any_of(
-#             contains_string("GDP per capita"),
-#             contains_string("GDP"),
-#         ),
-#     )
-#     assert_that(
-#         resp.business_insights,
-#         any_of(
-#             contains_string("United States"),
-#             contains_string("US"),
-#         ),
-#     )
-#     assert_that(
-#         resp.code_explanation,
-#         all_of(contains_string("https://vizro.readthedocs.io/en/stable/pages/user_guides/custom_charts/")),
-#     )
+import vizro.plotly.express as px
+from hamcrest import all_of, any_of, assert_that, contains_string, equal_to
+from vizro_ai import VizroAI
+
+vizro_ai = VizroAI()
+df = px.data.gapminder()
+
+
+def test_chart():
+    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,
+        all_of(contains_string("px.scatter")),
+    )
+    assert_that(
+        resp.code,
+        any_of(contains_string("x='continent'"), contains_string("x='gdpPercap'")),
+    )
+    assert_that(
+        resp.code,
+        any_of(contains_string("y='count'"), contains_string("y='gdpPercap'"), contains_string("y='continent'")),
+    )
+    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, return_elements=True)
+    assert_that(
+        resp.code,
+        all_of(contains_string("px.bar"), contains_string("x='year'")),
+    )
+    assert_that(
+        resp.code,
+        any_of(contains_string("y='gdpPercap'"), contains_string("y='total_gdp'")),
+    )
+    assert_that(
+        resp.business_insights,
+        any_of(
+            contains_string("GDP per capita"),
+            contains_string("GDP"),
+        ),
+    )
+    assert_that(
+        resp.business_insights,
+        any_of(
+            contains_string("United States"),
+            contains_string("US"),
+        ),
+    )
+    assert_that(
+        resp.code_explanation,
+        all_of(contains_string("https://vizro.readthedocs.io/en/stable/pages/user_guides/custom_charts/")),
+    )