momentry_core/scripts/test_identity_db.py

#!/opt/homebrew/bin/python3.11
"""
Test Identity Database Integration

Purpose: Verify identities table and reference_data JSONB storage

Usage:
    python3 scripts/test_identity_db.py
"""

import os
import sys
import json
import psycopg2
from datetime import datetime

DATABASE_URL = os.getenv("DATABASE_URL", "postgres://accusys@localhost:5432/momentry?options=-c%20search_path=dev")


def test_db_connection():
    """Test database connection"""
    print("🔧 Testing database connection...")
    try:
        conn = psycopg2.connect(DATABASE_URL)
        cur = conn.cursor()
        
        cur.execute("SELECT version();")
        version = cur.fetchone()[0]
        print(f"✅ Connected: {version}")
        
        cur.close()
        conn.close()
        return True
    except Exception as e:
        print(f"❌ Connection failed: {e}")
        return False


def test_identities_table(schema="dev"):
    """Test identities table structure"""
    print(f"\n🔧 Testing {schema}.identities table...")
    
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        cur.execute(f"""
            SELECT column_name, data_type 
            FROM information_schema.columns 
            WHERE table_schema = '{schema}' AND table_name = 'identities'
            ORDER BY ordinal_position;
        """)
        
        columns = cur.fetchall()
        print(f"✅ Table columns ({len(columns)}):")
        
        expected_columns = {
            "uuid": "uuid",
            "name": "character varying",
            "identity_type": "character varying",
            "source": "character varying",
            "status": "character varying",
            "face_embedding": "USER-DEFINED",
            "voice_embedding": "USER-DEFINED",
            "identity_embedding": "USER-DEFINED",
            "reference_data": "jsonb",
            "tmdb_id": "integer",
            "tmdb_profile": "text",
            "created_at": "timestamp with time zone",
            "updated_at": "timestamp with time zone",
        }
        
        for col_name, col_type in columns:
            expected_type = expected_columns.get(col_name)
            if expected_type and col_type != expected_type and col_type != "USER-DEFINED":
                print(f"  ⚠️ {col_name}: {col_type} (expected: {expected_type})")
            else:
                print(f"  ✅ {col_name}: {col_type}")
        
        return True
    except Exception as e:
        print(f"❌ Table check failed: {e}")
        return False
    finally:
        cur.close()
        conn.close()


def test_reference_data_storage(schema="dev"):
    """Test reference_data JSONB storage"""
    print("\n🔧 Testing reference_data JSONB storage...")
    
    test_data = {
        "face_embeddings": [
            {
                "embedding": [0.1] * 512,
                "source": "test",
                "image_url": "https://example.com/test.jpg",
                "angle": "frontal",
                "quality_score": 0.95,
                "created_at": datetime.now().isoformat(),
            }
        ],
        "identity_embeddings": [
            {
                "embedding": [0.2] * 768,
                "source": "logo_image",
                "image_url": "https://example.com/logo.png",
                "context": "brand_logo",
                "created_at": datetime.now().isoformat(),
            }
        ],
        "image_urls": ["https://example.com/test.jpg"],
    }
    
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        cur.execute(f"""
            INSERT INTO {schema}.identities (
                name, identity_type, source, status,
                face_embedding, identity_embedding, reference_data
            ) VALUES (
                'Test Identity', 'people', 'test', 'pending',
                %s, %s, %s
            )
            RETURNING uuid, reference_data;
        """, (
            "[" + ",".join(["0.1"] * 512) + "]",
            "[" + ",".join(["0.2"] * 768) + "]",
            json.dumps(test_data),
        ))
        
        uuid, stored_data = cur.fetchone()
        conn.commit()
        
        print(f"✅ Inserted test identity: {uuid}")
        
        stored_json = json.loads(stored_data) if isinstance(stored_data, str) else stored_data
        
        print("✅ Stored reference_data:")
        print(f"  - face_embeddings: {len(stored_json.get('face_embeddings', []))} items")
        print(f"  - identity_embeddings: {len(stored_json.get('identity_embeddings', []))} items")
        print(f"  - image_urls: {len(stored_json.get('image_urls', []))} items")
        
        cur.execute(f"DELETE FROM {schema}.identities WHERE name = 'Test Identity';")
        conn.commit()
        print("✅ Cleaned up test identity")
        
        return True
    except Exception as e:
        print(f"❌ Storage test failed: {e}")
        conn.rollback()
        return False
    finally:
        cur.close()
        conn.close()


def test_query_accusys_logo(schema="dev"):
    """Query Accusys Logo Identity"""
    print("\n🔧 Querying Accusys Logo Identity...")
    
    conn = psycopg2.connect(DATABASE_URL)
    cur = conn.cursor()
    
    try:
        cur.execute(f"""
            SELECT uuid, name, identity_type, source, status, reference_data
            FROM {schema}.identities
            WHERE name = 'Accusys Storage Logo';
        """)
        
        row = cur.fetchone()
        
        if row:
            uuid, name, identity_type, source, status, reference_data = row
            print("✅ Found Identity:")
            print(f"  - UUID: {uuid}")
            print(f"  - Name: {name}")
            print(f"  - Type: {identity_type}")
            print(f"  - Source: {source}")
            print(f"  - Status: {status}")
            
            ref_data = json.loads(reference_data) if isinstance(reference_data, str) else reference_data
            print("  - reference_data:")
            print(f"    - image_urls: {ref_data.get('image_urls', [])}")
            print(f"    - identity_embeddings: {len(ref_data.get('identity_embeddings', []))} items")
            
            return True
        else:
            print("⚠️ Accusys Logo Identity not found")
            return False
    except Exception as e:
        print(f"❌ Query failed: {e}")
        return False
    finally:
        cur.close()
        conn.close()


def main():
    print("=" * 60)
    print("Identity Database Integration Test")
    print("=" * 60)
    
    results = []
    
    results.append(("Database Connection", test_db_connection()))
    results.append(("Identities Table (dev)", test_identities_table("dev")))
    results.append(("Identities Table (public)", test_identities_table("public")))
    results.append(("reference_data Storage", test_reference_data_storage("dev")))
    results.append(("Accusys Logo Query", test_query_accusys_logo("dev")))
    
    print("\n" + "=" * 60)
    print("Test Results Summary")
    print("=" * 60)
    
    for test_name, passed in results:
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"{test_name}: {status}")
    
    all_passed = all(r[1] for r in results)
    
    print("\n" + "=" * 60)
    if all_passed:
        print("🎉 All tests passed!")
        sys.exit(0)
    else:
        print("❌ Some tests failed")
        sys.exit(1)


if __name__ == "__main__":
    main()