From c0d404799f24d912dccf2f08824b4a49f98925f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B0=81=E9=93=83=E8=A7=A3=E6=BB=A2?= <495950007@qq.com>
Date: Sat, 8 Jun 2024 23:04:57 +0800
Subject: [PATCH] Add unit tests for ExcelConverter and JSON validator
 functions

---
 tests/test_file_processing.py | 44 ++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/tests/test_file_processing.py b/tests/test_file_processing.py
index 442b199..2b8f63e 100644
--- a/tests/test_file_processing.py
+++ b/tests/test_file_processing.py
@@ -1 +1,43 @@
-# Test file processing
+import unittest
+import json
+import os
+import pandas as pd
+from ai_commons.file_processing.excel_parser import ExcelConverter
+from ai_commons.file_processing.json_validator import remove_keys_from_json
+
+class TestExcelConverter(unittest.TestCase):
+
+    def setUp(self):
+        # 创建一个示例Excel文件
+        self.test_file = 'test.xlsx'
+        df = pd.DataFrame({
+            'Timestamp': [pd.Timestamp('2024-04-07 20:50:30'), pd.Timestamp('2024-05-08 15:30:45')],
+            'User ID': ['U12345', 'U67890'],
+            'Gift': ['Flower', 'Chocolate']
+        })
+        df.to_excel(self.test_file, index=False)
+
+    def tearDown(self):
+        # 删除示例Excel文件
+        if os.path.exists(self.test_file):
+            os.remove(self.test_file)
+
+    def test_to_json(self):
+        converter = ExcelConverter(self.test_file)
+        json_output = converter.to_json(chunk_size=2)
+        expected_output = [
+            '[{"Timestamp": "2024-04-07 20:50:30", "User ID": "U12345", "Gift": "Flower"}, {"Timestamp": "2024-05-08 15:30:45", "User ID": "U67890", "Gift": "Chocolate"}]'
+        ]
+        self.assertEqual(json_output, expected_output)
+
+class TestJsonValidator(unittest.TestCase):
+
+    def test_remove_keys_from_json(self):
+        json_string = '[{"Timestamp": "2024-04-07 20:50:30", "User ID": "U12345", "Gift": "Flower"}, {"Timestamp": "2024-05-08 15:30:45", "User ID": "U67890", "Gift": "Chocolate"}]'
+        keys_to_remove = [["Timestamp"], ["User ID"]]
+        modified_json_string = remove_keys_from_json(json_string, keys_to_remove)
+        expected_output = '[{"Gift": "Flower"}, {"Gift": "Chocolate"}]'
+        self.assertEqual(modified_json_string, expected_output)
+
+if __name__ == '__main__':
+    unittest.main()