def test_save_for_entity_and_edge(kb: KB, apple, google):
assert apple == kb.save(apple)
assert google == kb.save(google)
kb.reindex()
assert 2 == len(kb)
assert apple == kb.get_node(apple.key)
kb.connect(start=apple, verb="IS_A", end=apple)
kb.reindex()
assert kb.info()["graph"] == {
"nodes": 2,
"edges": 1,
}
assert 2 == len(kb.get_edges(node_key=apple))
assert 1 == len(kb.get_edges(node_key=apple, direction=Direction.incoming))
assert 2 == len(kb.get_edges(node_key=apple, verb="IS_A"))
assert 1 == len(kb.get_edges(node_key=apple, verb="IS_A", limit=1))
assert 0 == len(kb.get_edges(node_key=apple, verb="IS_NOT"))
assert apple.key == kb.get_neighbors(apple).neighbors[0].key
assert ([] == kb.get_neighbors(apple,
verb="IS_NOT",
direction=Direction.outgoing).neighbors)
kb.save(Edge(start=apple, verb="POINTS_NO_WHERE", end="INVALID|THING"))
kb.save(Edge(start=apple, verb="POINTS_NO_WHERE", end=google))
kb.reindex()
assert kb.info()["graph"] == {
"nodes": 2,
"edges": 3,
}
t = T().all_nodes(passthru=True)
response = kb.search(q="a", traversal=t)
assert 3 == len(response.nodes)
kb.remove_node(apple.key)
kb.reindex()
assert kb.info()["graph"] == {
"nodes": 1,
"edges": 3,
}
kb.clean_edges()
assert kb.info()["graph"] == {
"nodes": 1,
"edges": 0,
}
data = response.dict()
compare = SearchResponse(**data)
assert compare.nodes == response.nodes
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"