From 882efe1dcd872e7b7c4700c31bce3807c66db887 Mon Sep 17 00:00:00 2001
From: Lingyi Zhang <lingyi_zhang@mckinsey.com>
Date: Fri, 21 Jun 2024 21:24:19 -0400
Subject: [PATCH] update unit tests

---
 vizro-ai/hatch.toml                                            | 3 ++-
 .../tests/unit/vizro-ai/components/test_chart_selection.py     | 2 +-
 .../tests/unit/vizro-ai/components/test_code_validation.py     | 2 +-
 .../tests/unit/vizro-ai/components/test_custom_chart_wrap.py   | 2 +-
 .../tests/unit/vizro-ai/components/test_dataframe_craft.py     | 2 +-
 vizro-ai/tests/unit/vizro-ai/components/test_explanation.py    | 2 +-
 vizro-ai/tests/unit/vizro-ai/components/test_visual_code.py    | 2 +-
 7 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/vizro-ai/hatch.toml b/vizro-ai/hatch.toml
index e82532d14..66b69e519 100644
--- a/vizro-ai/hatch.toml
+++ b/vizro-ai/hatch.toml
@@ -16,7 +16,8 @@ dependencies = [
   "toml",
   "nbformat>=4.2.0",
   "pyhamcrest",
-  "jupyter"
+  "jupyter",
+  "langchain_community"
 ]
 
 [envs.default.env-vars]
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_chart_selection.py b/vizro-ai/tests/unit/vizro-ai/components/test_chart_selection.py
index 9f6246055..5a83bc2d4 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_chart_selection.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_chart_selection.py
@@ -1,6 +1,6 @@
 import pandas as pd
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetChartSelection
 
 
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_code_validation.py b/vizro-ai/tests/unit/vizro-ai/components/test_code_validation.py
index 4172e1c3a..1c8f42108 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_code_validation.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_code_validation.py
@@ -1,5 +1,5 @@
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetDebugger
 
 
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_custom_chart_wrap.py b/vizro-ai/tests/unit/vizro-ai/components/test_custom_chart_wrap.py
index 882d33422..e1e055c50 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_custom_chart_wrap.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_custom_chart_wrap.py
@@ -1,5 +1,5 @@
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetCustomChart
 
 
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_dataframe_craft.py b/vizro-ai/tests/unit/vizro-ai/components/test_dataframe_craft.py
index ff605ee1c..7e40ea528 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_dataframe_craft.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_dataframe_craft.py
@@ -2,7 +2,7 @@
 
 import pandas as pd
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetDataFrameCraft
 
 
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_explanation.py b/vizro-ai/tests/unit/vizro-ai/components/test_explanation.py
index 51a06c1f4..c6eb35c87 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_explanation.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_explanation.py
@@ -1,5 +1,5 @@
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetCodeExplanation
 
 
diff --git a/vizro-ai/tests/unit/vizro-ai/components/test_visual_code.py b/vizro-ai/tests/unit/vizro-ai/components/test_visual_code.py
index f9132a954..0ba9f19a7 100644
--- a/vizro-ai/tests/unit/vizro-ai/components/test_visual_code.py
+++ b/vizro-ai/tests/unit/vizro-ai/components/test_visual_code.py
@@ -1,5 +1,5 @@
 import pytest
-from langchain.llms.fake import FakeListLLM
+from langchain_community.llms.fake import FakeListLLM
 from vizro_ai.components import GetVisualCode