test_json.py 1.0 KB

1234567891011121314151617181920212223242526272829303132
  1. import hashlib
  2. from unittest.mock import patch
  3. from langchain.docstore.document import Document
  4. from langchain.document_loaders.json_loader import \
  5. JSONLoader as LangchainJSONLoader
  6. from embedchain.loaders.json import JSONLoader
  7. def test_load_data():
  8. mock_document = [
  9. Document(page_content="content1", metadata={"seq_num": 1}),
  10. Document(page_content="content2", metadata={"seq_num": 2}),
  11. ]
  12. with patch.object(LangchainJSONLoader, "load", return_value=mock_document):
  13. content = "temp.json"
  14. result = JSONLoader.load_data(content)
  15. assert "doc_id" in result
  16. assert "data" in result
  17. expected_data = [
  18. {"content": "content1", "meta_data": {"url": content, "row": 1}},
  19. {"content": "content2", "meta_data": {"url": content, "row": 2}},
  20. ]
  21. assert result["data"] == expected_data
  22. expected_doc_id = hashlib.sha256((content + ", ".join(["content1", "content2"])).encode()).hexdigest()
  23. assert result["doc_id"] == expected_doc_id