33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
import embeddings
|
|
|
|
|
|
class TextExtractionTests(unittest.TestCase):
|
|
def write_text(self, suffix, content):
|
|
handle = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
|
|
handle.close()
|
|
self.addCleanup(lambda: os.path.exists(handle.name) and os.unlink(handle.name))
|
|
with open(handle.name, "w", encoding="utf-8") as stream:
|
|
stream.write(content)
|
|
return handle.name
|
|
|
|
def test_extracts_utf8_markdown(self):
|
|
path = self.write_text(".md", "# 退款规则\n\n七日内可申请退款。")
|
|
self.assertEqual(embeddings.extract_text(path, ".md"), "# 退款规则\n\n七日内可申请退款。")
|
|
|
|
def test_extracts_utf8_text(self):
|
|
path = self.write_text(".txt", "客服热线：400-123-4567")
|
|
self.assertEqual(embeddings.extract_text(path, ".txt"), "客服热线：400-123-4567")
|
|
|
|
def test_rejects_unsupported_extension(self):
|
|
path = self.write_text(".csv", "not supported")
|
|
with self.assertRaises(ValueError):
|
|
embeddings.extract_text(path, ".csv")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|