embedchain

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

4
import pytest
5

6
from embedchain.loaders.youtube_video import YoutubeVideoLoader
7

8

9
@pytest.fixture
10
def youtube_video_loader():
11
    return YoutubeVideoLoader()
12

13

14
def test_load_data(youtube_video_loader):
15
    video_url = "https://www.youtube.com/watch?v=VIDEO_ID"
16
    mock_loader = Mock()
17
    mock_page_content = "This is a YouTube video content."
18
    mock_loader.load.return_value = [
19
        MagicMock(
20
            page_content=mock_page_content,
21
            metadata={"url": video_url, "title": "Test Video"},
22
        )
23
    ]
24

25
    with patch("embedchain.loaders.youtube_video.YoutubeLoader.from_youtube_url", return_value=mock_loader):
26
        result = youtube_video_loader.load_data(video_url)
27

28
    expected_doc_id = hashlib.sha256((mock_page_content + video_url).encode()).hexdigest()
29

30
    assert result["doc_id"] == expected_doc_id
31

32
    expected_data = [
33
        {
34
            "content": "This is a YouTube video content.",
35
            "meta_data": {"url": video_url, "title": "Test Video"},
36
        }
37
    ]
38

39
    assert result["data"] == expected_data
40

41

42
def test_load_data_with_empty_doc(youtube_video_loader):
43
    video_url = "https://www.youtube.com/watch?v=VIDEO_ID"
44
    mock_loader = Mock()
45
    mock_loader.load.return_value = []
46

47
    with patch("embedchain.loaders.youtube_video.YoutubeLoader.from_youtube_url", return_value=mock_loader):
48
        with pytest.raises(ValueError):
49
            youtube_video_loader.load_data(video_url)
50