embedchain

1
import os
2
from unittest.mock import MagicMock
3

4
import pytest
5
from dropbox.files import FileMetadata
6

7
from embedchain.loaders.dropbox import DropboxLoader
8

9

10
@pytest.fixture
11
def setup_dropbox_loader(mocker):
12
    mock_dropbox = mocker.patch("dropbox.Dropbox")
13
    mock_dbx = mocker.MagicMock()
14
    mock_dropbox.return_value = mock_dbx
15

16
    os.environ["DROPBOX_ACCESS_TOKEN"] = "test_token"
17
    loader = DropboxLoader()
18

19
    yield loader, mock_dbx
20

21
    if "DROPBOX_ACCESS_TOKEN" in os.environ:
22
        del os.environ["DROPBOX_ACCESS_TOKEN"]
23

24

25
def test_initialization(setup_dropbox_loader):
26
    """Test initialization of DropboxLoader."""
27
    loader, _ = setup_dropbox_loader
28
    assert loader is not None
29

30

31
def test_download_folder(setup_dropbox_loader, mocker):
32
    """Test downloading a folder."""
33
    loader, mock_dbx = setup_dropbox_loader
34
    mocker.patch("os.makedirs")
35
    mocker.patch("os.path.join", return_value="mock/path")
36

37
    mock_file_metadata = mocker.MagicMock(spec=FileMetadata)
38
    mock_dbx.files_list_folder.return_value.entries = [mock_file_metadata]
39

40
    entries = loader._download_folder("path/to/folder", "local_root")
41
    assert entries is not None
42

43

44
def test_generate_dir_id_from_all_paths(setup_dropbox_loader, mocker):
45
    """Test directory ID generation."""
46
    loader, mock_dbx = setup_dropbox_loader
47
    mock_file_metadata = mocker.MagicMock(spec=FileMetadata, name="file.txt")
48
    mock_dbx.files_list_folder.return_value.entries = [mock_file_metadata]
49

50
    dir_id = loader._generate_dir_id_from_all_paths("path/to/folder")
51
    assert dir_id is not None
52
    assert len(dir_id) == 64
53

54

55
def test_clean_directory(setup_dropbox_loader, mocker):
56
    """Test cleaning up a directory."""
57
    loader, _ = setup_dropbox_loader
58
    mocker.patch("os.listdir", return_value=["file1", "file2"])
59
    mocker.patch("os.remove")
60
    mocker.patch("os.rmdir")
61

62
    loader._clean_directory("path/to/folder")
63

64

65
def test_load_data(mocker, setup_dropbox_loader, tmp_path):
66
    loader = setup_dropbox_loader[0]
67

68
    mock_file_metadata = MagicMock(spec=FileMetadata, name="file.txt")
69
    mocker.patch.object(loader.dbx, "files_list_folder", return_value=MagicMock(entries=[mock_file_metadata]))
70
    mocker.patch.object(loader.dbx, "files_download_to_file")
71

72
    # Mock DirectoryLoader
73
    mock_data = {"data": "test_data"}
74
    mocker.patch("embedchain.loaders.directory_loader.DirectoryLoader.load_data", return_value=mock_data)
75

76
    test_dir = tmp_path / "dropbox_test"
77
    test_dir.mkdir()
78
    test_file = test_dir / "file.txt"
79
    test_file.write_text("dummy content")
80
    mocker.patch.object(loader, "_generate_dir_id_from_all_paths", return_value=str(test_dir))
81

82
    result = loader.load_data("path/to/folder")
83

84
    assert result == {"doc_id": mocker.ANY, "data": "test_data"}
85
    loader.dbx.files_list_folder.assert_called_once_with("path/to/folder")
86