Spaces:
Sleeping
Sleeping
File size: 1,426 Bytes
cfd3735 |
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 |
import unittest
from langchain.document_loaders.duckdb_loader import DuckDBLoader
try:
import duckdb # noqa: F401
duckdb_installed = True
except ImportError:
duckdb_installed = False
@unittest.skipIf(not duckdb_installed, "duckdb not installed")
def test_duckdb_loader_no_options() -> None:
"""Test DuckDB loader."""
loader = DuckDBLoader("SELECT 1 AS a, 2 AS b")
docs = loader.load()
assert len(docs) == 1
assert docs[0].page_content == "a: 1\nb: 2"
assert docs[0].metadata == {}
@unittest.skipIf(not duckdb_installed, "duckdb not installed")
def test_duckdb_loader_page_content_columns() -> None:
"""Test DuckDB loader."""
loader = DuckDBLoader(
"SELECT 1 AS a, 2 AS b UNION SELECT 3 AS a, 4 AS b",
page_content_columns=["a"],
)
docs = loader.load()
assert len(docs) == 2
assert docs[0].page_content == "a: 1"
assert docs[0].metadata == {}
assert docs[1].page_content == "a: 3"
assert docs[1].metadata == {}
@unittest.skipIf(not duckdb_installed, "duckdb not installed")
def test_duckdb_loader_metadata_columns() -> None:
"""Test DuckDB loader."""
loader = DuckDBLoader(
"SELECT 1 AS a, 2 AS b",
page_content_columns=["a"],
metadata_columns=["b"],
)
docs = loader.load()
assert len(docs) == 1
assert docs[0].page_content == "a: 1"
assert docs[0].metadata == {"b": 2}
|