test_checkpoint.py

"""
Test script for checkpoint functionality.

Tests checkpoint save and resume by processing opinions in batches.
"""

import asyncio
import json
import os
import sys
from pathlib import Path

# Add project root to path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from dotenv import load_dotenv
from src.api_client import fetch_all_opinions

load_dotenv()
TOKEN = os.getenv('COURTLISTENER_TOKEN')

if not TOKEN:
    print("Error: COURTLISTENER_TOKEN not found in .env")
    sys.exit(1)

# Load some opinion IDs from the preprocessed data
opinions_path = PROJECT_ROOT / 'data' / 'opinions.json'
with open(opinions_path, 'r', encoding='utf-8') as f:
    opinions = json.load(f)

# Take first 250 opinion IDs for testing
test_ids = [op['id'] for op in opinions[:250] if op.get('id')]

print(f"Testing checkpoint functionality with {len(test_ids)} opinions...")
print(f"This will take approximately {len(test_ids) / 5 / 60:.1f} minutes")
print("\nCheckpoint directory: .checkpoints/")
print("You can interrupt this process (Ctrl+C) and restart to test resume functionality")
print("\n" + "="*60)

checkpoint_dir = PROJECT_ROOT / '.checkpoints'

async def run_test():
    results = await fetch_all_opinions(
        test_ids,
        TOKEN,
        batch_size=100,
        checkpoint_dir=checkpoint_dir
    )
    return results

try:
    results = asyncio.run(run_test())
    print(f"\n[OK] Completed successfully!")
    print(f"Fetched {len(results)}/{len(test_ids)} opinions")
    
    # Show sample result
    if results:
        sample_id = list(results.keys())[0]
        sample = results[sample_id]
        html_length = len(sample.get('html_with_citations', ''))
        print(f"\nSample opinion {sample_id}: {html_length} chars of HTML")
        
except KeyboardInterrupt:
    print("\n\n[INTERRUPTED] Process stopped by user")
    print(f"Checkpoint saved in: {checkpoint_dir}")
    print("Re-run this script to resume from checkpoint")
    sys.exit(0)