File size: 2,271 Bytes
220f34e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
468ae07
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import pytest
from unittest.mock import patch, MagicMock
from connectors.mongodb_handler import MongodbHandler, list_mongodb_collections
import pandas as pd

@patch('connectors.mongodb_handler.MongoClient')
def test_list_mongodb_collections(mock_client_class):
    mock_client = MagicMock()
    mock_db = MagicMock()
    mock_db.list_collection_names.return_value = ["users", "orders"]
    mock_client.__getitem__.return_value = mock_db
    mock_client_class.return_value = mock_client

    collections = list_mongodb_collections("localhost", 27017, "testdb", "user", "pass")
    assert collections == ["users", "orders"]

@patch('connectors.mongodb_handler.MongoClient')
def test_fetch_data_mongodb(mock_client_class):
    mock_client = MagicMock()
    mock_db = MagicMock()
    mock_collection = MagicMock()
    
    # Mock finding docs
    mock_cursor = MagicMock()
    mock_cursor.limit.return_value = [
        {"_id": "123", "user": {"name": "Alice", "age": 25}},
        {"_id": "456", "user": {"name": "Bob", "age": 30}}
    ]
    mock_collection.find.return_value = mock_cursor
    
    mock_db.__getitem__.return_value = mock_collection
    mock_client.__getitem__.return_value = mock_db
    mock_client_class.return_value = mock_client

    handler = MongodbHandler()
    df = handler.fetch_data("localhost", 27017, "testdb", "user", "pass", "users")
    
    assert isinstance(df, pd.DataFrame)
    assert "user.name" in df.columns
    assert "user.age" in df.columns
    assert df.iloc[0]["user.name"] == "Alice"
    assert df.iloc[1]["user.name"] == "Bob"

@patch('connectors.mongodb_handler.MongoClient')
def test_fetch_data_respects_limit(mock_client_class):
    mock_client = MagicMock()
    mock_db = MagicMock()
    mock_collection = MagicMock()
    mock_cursor = MagicMock()
    
    mock_cursor.limit.return_value = []
    mock_collection.find.return_value = mock_cursor
    
    mock_db.__getitem__.return_value = mock_collection
    mock_client.__getitem__.return_value = mock_db
    mock_client_class.return_value = mock_client

    handler = MongodbHandler()
    handler.fetch_data("localhost", 27017, "testdb", "user", "pass", "users", limit=5)
    
    mock_collection.find.assert_called_once()
    mock_cursor.limit.assert_called_once_with(5)