From cd3bf4294e13131a9d94fdcc1e6ae6bd9ae7b02d Mon Sep 17 00:00:00 2001
From: cFireworks <xwtfd@outlook.com>
Date: Thu, 8 Aug 2024 11:23:39 +0800
Subject: [PATCH] fix: Support Chinese characters in prompt generation stage
 (#1168) Update dataframe_serializer.py add test case to #1168

---
 .../helpers/test_dataframe_serializer.py      | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 tests/unit_tests/helpers/test_dataframe_serializer.py

diff --git a/tests/unit_tests/helpers/test_dataframe_serializer.py b/tests/unit_tests/helpers/test_dataframe_serializer.py
new file mode 100644
index 000000000..3cf64b7df
--- /dev/null
+++ b/tests/unit_tests/helpers/test_dataframe_serializer.py
@@ -0,0 +1,30 @@
+import unittest
+
+import pandas as pd
+
+from pandasai.connectors import PandasConnector
+from pandasai.helpers.dataframe_serializer import (
+    DataframeSerializer,
+    DataframeSerializerType,
+)
+
+
+class TestDataframeSerializer(unittest.TestCase):
+    def setUp(self):
+        self.serializer = DataframeSerializer()
+
+    def test_convert_df_to_yml(self):
+        # Test convert df to yml
+        data = {"name": ["en_name", "中文_名称"]}
+        connector = PandasConnector(
+            {"original_df": pd.DataFrame(data)},
+            name="en_table_name",
+            description="中文_描述",
+            field_descriptions={k: k for k in data},
+        )
+        result = self.serializer.serialize(
+            connector,
+            type_=DataframeSerializerType.YML,
+            extras={"index": 0, "type": "pd.Dataframe"},
+        )
+        self.assertIn("中文_描述", result)