tests/test_graph_visual.py

"""
Test the PyVis graph visualiser.
Loads a sample, builds graph, picks a fraud account, and renders to HTML.
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import warnings
warnings.filterwarnings('ignore')

from src.data_loader import get_processed_data
import src.graph_builder as gb
from src.visualiser.pyvis_graph import build_pyvis_graph, save_pyvis_html

OUTPUT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'graph_test_output.html')

print("Loading data...")
df, nf = get_processed_data()

# Use a 1000-row sample for faster testing
sample_df = df.sample(n=min(1000, len(df)), random_state=42).copy()
print(f"Sample: {len(sample_df)} rows")

print("Building graph from sample...")
G = gb.build_graph(sample_df)
G = gb.attach_node_features(G, nf)
print(f"Sample graph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

print("Computing PageRank on sample graph...")
pagerank_scores = gb.compute_pagerank(G)

print("Computing Louvain on sample graph...")
import community as community_louvain
G_undirected = G.to_undirected()
louvain_partition = community_louvain.best_partition(G_undirected, weight='amount')

# Pick a confirmed fraud account
fraud_accounts = set(df[df['is_laundering'] == 1]['source'].values)
fraud_in_sample = fraud_accounts & set(G.nodes())
if fraud_in_sample:
    center_node = list(fraud_in_sample)[0]
else:
    center_node = list(G.nodes())[0]
print(f"Center node: {center_node} (fraud={center_node in fraud_accounts})")

# Extract subgraph
sub_G = gb.get_subgraph(G, center_node, hops=2, max_nodes=60)
print(f"Subgraph: {sub_G.number_of_nodes()} nodes, {sub_G.number_of_edges()} edges")

# Build PyVis graph
print("Building PyVis graph...")
net = build_pyvis_graph(
    subgraph=sub_G,
    df=sample_df,
    center_node=center_node,
    fraud_accounts=fraud_accounts,
    pagerank_scores=pagerank_scores,
    louvain_partition=louvain_partition,
)

# Save HTML
save_pyvis_html(net, OUTPUT_PATH)
print(f"\n✅ Graph built successfully with {sub_G.number_of_nodes()} nodes and {sub_G.number_of_edges()} edges")
print(f"HTML saved to: {OUTPUT_PATH}")

# Verify file exists and has content
import os
size = os.path.getsize(OUTPUT_PATH)
print(f"HTML file size: {size:,} bytes")