DB-GPT/dbgpt/rag/knowledge/tests/test_csv.py
2024-01-14 21:01:37 +08:00

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