!/usr/bin/env python3

"""Quick test of the knowledge graph module."""
import json
import sys
sys.path.insert(0, '.')
from graph import (
add_entity, add_relation, get_entity, search_entities,
explore, health, merge_entities, init_db, _rebuild_index,
VALID_RELATION_TYPES
)

Ensure schema is initialized

init_db()

Clean up any test data from previous runs

from graph import DB_PATH
import sqlite3
conn = sqlite3.connect(DB_PATH)
conn.execute('DELETE FROM relations')
conn.execute('DELETE FROM entities')
conn.commit()
conn.close()

Rebuild empty index

_rebuild_index()

Test health

print('=== HEALTH ===')
h = health()
print(json.dumps(h, indent=2, default=str))

Test adding entities

print('\n=== ADD ENTITIES ===')
e1 = add_entity('person', 'Matt Hoffmann',
'Matt lives in Green Bay, WI. Husband to Aundrea, father to Sullivan and Harper.',
tags=['green_bay', 'family'], confidence=1.0)
print(f'Entity 1: {e1}')

e2 = add_entity('person', 'Aundrea Hoffmann',
'Aundrea is Matts wife.',
tags=['family'])
print(f'Entity 2: {e2}')

e3 = add_entity('place', 'Green Bay',
'Green Bay, Wisconsin. Home of the Hoffmann family.',
tags=['wisconsin', 'home'])
print(f'Entity 3: {e3}')

Test adding relations

print('\n=== ADD RELATIONS ===')
r1 = add_relation(e1, e2, 'spouse_of', confidence=1.0)
print(f'Relation 1: {r1}')
r2 = add_relation(e1, e3, 'lives_in', confidence=1.0)
print(f'Relation 2: {r2}')

Test get_entity

print('\n=== GET ENTITY ===')
entity = get_entity(e1)
print(f'Entity: {entity["subject"]} ({entity["type"]})')
print(f'Relations out: {len(entity["relations_out"])}')
print(f'Relations in: {len(entity["relations_in"])}')
for rel in entity['relations_out']:
print(f' -> {rel["relation_type"]} -> {rel["target_subject"]}')

Test explore

print('\n=== EXPLORE ===')
ex = explore(e1, depth=1)
print(f'Root: {ex["root"]["subject"]}')
for nid, node in ex.get('nodes', {}).items():
print(f' Neighbor: {node["subject"]} ({node["type"]})')

Test semantic search

print('\n=== SEARCH ===')
results = search_entities('Matts family', top_k=3)
for r in results:
print(f' [{r["similarity"]:.3f}] {r["subject"]} ({r["type"]})')

Test dedup

print('\n=== DEDUP ===')
e4 = add_entity('person', 'matt hoffmann',
'Lives in WI with his family.',
tags=['green_bay'])
print(f'Dedup returned: {e4} (should equal {e1})')
assert e4 == e1, 'Dedup should return existing entity ID'
print(' ✓ Dedup working')

Test vague relation rejection

print('\n=== VAGUE RELATION REJECTION ===')
try:
add_relation(e1, e3, 'related_to')
print(' FAILED: Should have rejected "related_to"')
except ValueError as e:
print(f' ✓ Rejected: {e}')

print('\n=== All tests passed! ===')