embedchain

1
import pytest
2

3
from embedchain.loaders.google_drive import GoogleDriveLoader
4

5

6
@pytest.fixture
7
def google_drive_folder_loader():
8
    return GoogleDriveLoader()
9

10

11
def test_load_data_invalid_drive_url(google_drive_folder_loader):
12
    mock_invalid_drive_url = "https://example.com"
13
    with pytest.raises(
14
        ValueError,
15
        match="The url provided https://example.com does not match a google drive folder url. Example "
16
        "drive url: https://drive.google.com/drive/u/0/folders/xxxx",
17
    ):
18
        google_drive_folder_loader.load_data(mock_invalid_drive_url)
19

20

21
@pytest.mark.skip(reason="This test won't work unless google api credentials are properly setup.")
22
def test_load_data_incorrect_drive_url(google_drive_folder_loader):
23
    mock_invalid_drive_url = "https://drive.google.com/drive/u/0/folders/xxxx"
24
    with pytest.raises(
25
        FileNotFoundError, match="Unable to locate folder or files, check provided drive URL and try again"
26
    ):
27
        google_drive_folder_loader.load_data(mock_invalid_drive_url)
28

29

30
@pytest.mark.skip(reason="This test won't work unless google api credentials are properly setup.")
31
def test_load_data(google_drive_folder_loader):
32
    mock_valid_url = "YOUR_VALID_URL"
33
    result = google_drive_folder_loader.load_data(mock_valid_url)
34
    assert "doc_id" in result
35
    assert "data" in result
36
    assert "content" in result["data"][0]
37
    assert "meta_data" in result["data"][0]
38