embedchain

1
import hashlib
2
import os
3
from unittest.mock import Mock, patch
4

5
import pytest
6

7
from embedchain.loaders.notion import NotionLoader
8

9

10
@pytest.fixture
11
def notion_loader():
12
    with patch.dict(os.environ, {"NOTION_INTEGRATION_TOKEN": "test_notion_token"}):
13
        yield NotionLoader()
14

15

16
def test_load_data(notion_loader):
17
    source = "https://www.notion.so/Test-Page-1234567890abcdef1234567890abcdef"
18
    mock_text = "This is a test page."
19
    expected_doc_id = hashlib.sha256((mock_text + source).encode()).hexdigest()
20
    expected_data = [
21
        {
22
            "content": mock_text,
23
            "meta_data": {"url": "notion-12345678-90ab-cdef-1234-567890abcdef"},  # formatted_id
24
        }
25
    ]
26

27
    mock_page = Mock()
28
    mock_page.text = mock_text
29
    mock_documents = [mock_page]
30

31
    with patch("embedchain.loaders.notion.NotionPageLoader") as mock_reader:
32
        mock_reader.return_value.load_data.return_value = mock_documents
33
        result = notion_loader.load_data(source)
34

35
    assert result["doc_id"] == expected_doc_id
36
    assert result["data"] == expected_data
37