mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-07-26 21:37:40 +00:00
33 lines
948 B
Python
33 lines
948 B
Python
from unittest.mock import MagicMock, mock_open, patch
|
|
|
|
import pytest
|
|
|
|
from dbgpt.rag.knowledge.csv import CSVKnowledge
|
|
|
|
MOCK_CSV_DATA = "id,name,age\n1,John Doe,30\n2,Jane Smith,25\n3,Bob Johnson,40"
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_file_open():
|
|
with patch("builtins.open", mock_open(read_data=MOCK_CSV_DATA)) as mock_file:
|
|
yield mock_file
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_csv_dict_reader():
|
|
with patch("csv.DictReader", MagicMock()) as mock_csv:
|
|
mock_csv.return_value = iter(
|
|
[
|
|
{"id": "1", "name": "John Doe", "age": "30"},
|
|
{"id": "2", "name": "Jane Smith", "age": "25"},
|
|
{"id": "3", "name": "Bob Johnson", "age": "40"},
|
|
]
|
|
)
|
|
yield mock_csv
|
|
|
|
|
|
def test_load_from_csv(mock_file_open, mock_csv_dict_reader):
|
|
knowledge = CSVKnowledge(file_path="test_data.csv", source_column="name")
|
|
documents = knowledge._load()
|
|
assert len(documents) == 3
|