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()
|