Spaces:
Sleeping
Sleeping
File size: 1,811 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 |
"""Test in memory docstore."""
import pytest
from langchain.docstore.document import Document
from langchain.docstore.in_memory import InMemoryDocstore
def test_document_found() -> None:
"""Test document found."""
_dict = {"foo": Document(page_content="bar")}
docstore = InMemoryDocstore(_dict)
output = docstore.search("foo")
assert isinstance(output, Document)
assert output.page_content == "bar"
def test_document_not_found() -> None:
"""Test when document is not found."""
_dict = {"foo": Document(page_content="bar")}
docstore = InMemoryDocstore(_dict)
output = docstore.search("bar")
assert output == "ID bar not found."
def test_adding_document() -> None:
"""Test that documents are added correctly."""
_dict = {"foo": Document(page_content="bar")}
docstore = InMemoryDocstore(_dict)
new_dict = {"bar": Document(page_content="foo")}
docstore.add(new_dict)
# Test that you can find new document.
foo_output = docstore.search("bar")
assert isinstance(foo_output, Document)
assert foo_output.page_content == "foo"
# Test that old document is the same.
bar_output = docstore.search("foo")
assert isinstance(bar_output, Document)
assert bar_output.page_content == "bar"
def test_adding_document_already_exists() -> None:
"""Test that error is raised if document id already exists."""
_dict = {"foo": Document(page_content="bar")}
docstore = InMemoryDocstore(_dict)
new_dict = {"foo": Document(page_content="foo")}
# Test that error is raised.
with pytest.raises(ValueError):
docstore.add(new_dict)
# Test that old document is the same.
bar_output = docstore.search("foo")
assert isinstance(bar_output, Document)
assert bar_output.page_content == "bar"
|