def test_search_with_results(kb: KB, apple, google):
kb.save_node(apple)
kb.save_node(google)
kb.reindex()
# default (all nodes, no filter, etc.)
response = kb.search()
assert [apple, google] == response.nodes
# offset = 1, skips 1 node
response = kb.search(offset=1)
assert [google] == response.nodes
# limit = 1
response = kb.search(limit=1)
assert [apple] == response.nodes
# prefix
response = kb.search(q="a")
assert [apple] == response.nodes
# keys
response = kb.search(keys=["Apple, Inc.|COMPANY"])
assert [apple] == response.nodes
# keys
response = kb.search(keys=[apple.key, apple.key, "junk"])
assert [apple] == response.nodes
# labels
response = kb.search(labels=["COMPANY"])
assert 2 == len(response.nodes)
# keys + labels
response = kb.search(keys=["Apple, Inc.|COMPANY"], labels=["COMPANY"])
assert [apple] == response.nodes
# dict
assert response.dict() == {
"nodes": [{
"data": None,
"headquarters": {
"city": "Cupertino",
"data": None,
"key": "1",
"label": "LOCATION",
},
"key": "Apple, Inc.|COMPANY",
"label": "COMPANY",
"name": "Apple, Inc.",
"synonyms": ("Apple", "AAPL"),
}],
"trails": [{
"end": "Apple, Inc.|COMPANY",
"hops": [],
"start": "Apple, Inc.|COMPANY",
}],
}
def test_search_with_just_text(kb: KB, apple, google):
kb.save_node(apple)
kb.save_node(google)
kb.reindex()
response = kb.search("ap")
assert 1 == len(response)
assert [apple] == response.nodes
assert apple == response[0]
assert [apple] == list(response)
def test_save_load_sync(root, kb: KB, apple):
def check():
assert (kb.parse("AAPL")).spans[0].entity == apple
assert (kb.parse("Apple, Inc.")).spans[0].entity == apple
assert (kb.parse("Apple,Inc.")).spans[0].entity == apple
with kb.transact():
kb.save_node(apple)
kb.reindex()
check()
kb = KB(root=root)
check()
kb.reload()
check()
def test_save_entity(kb: KB, apple, apple_records):
kb.save_node(apple)
kb.save_node(apple_records)
assert {apple, apple_records} == set(kb)
kb.reindex()
# parse functions
assert (kb.parse("AAPL")).spans[0].entity == apple
assert (kb.parse("Apple, Inc.")).spans[0].entity == apple
assert (kb.parse("Apple Computers")).spans[0].text == "Apple"
assert (kb.parse("Apple Records")).spans[0].entity == apple_records
assert 2 == len((kb.parse("Apple")).spans)
# find functions
assert 2 == len(kb.find("apple"))
assert kb.find_one("apple") is None # 2 results cause no return
assert kb.find_one("AAPL").name == "Apple, Inc."
# should reset the terms
apple2 = apple.copy(update=dict(synonyms=("Apple", "Apple Computers")))
kb.save_node(apple2)
kb.reindex()
assert not (kb.parse("AAPL")).spans
assert (kb.parse("Apple, Inc.")).spans[0].entity == apple2
assert (kb.parse("Apple Computers")).spans[0].entity == apple2
assert (kb.parse("Apple Computers")).spans[0].text == "Apple Computers"
assert 2 == len((kb.parse("Apple")).spans)
kb.remove_node(apple2)
kb.reindex()
assert 1 == len((kb.parse("Apple")).spans)
assert 1 == len((kb.parse("Apple Computers")).spans)
assert (kb.parse("Apple Computers")).spans[0].text == "Apple"
