test_backend.py

#!/usr/bin/env python3
"""
Test Backend for RAG Chatbot Knowledge Base
Tests the ChromaDB vector database logic directly, bypassing the Streamlit UI.

Run with: python test_backend.py
"""

import os
import shutil
import sys

TEST_DB_PATH = "./test_chroma_db"

def setup_test_db():
    if os.path.exists(TEST_DB_PATH):
        shutil.rmtree(TEST_DB_PATH)

import vector_db
vector_db.PERSIST_DIRECTORY = TEST_DB_PATH
vector_db._client = None
vector_db._collection = None

def cleanup_test_db():
    vector_db._client = None
    vector_db._collection = None
    if os.path.exists(TEST_DB_PATH):
        shutil.rmtree(TEST_DB_PATH)

def test_case_1_fresh_ingestion():
    print("\n" + "="*60)
    print("TEST CASE 1: Fresh Ingestion & Metadata")
    print("="*60)
    
    setup_test_db()
    
    chunk_count = vector_db.ingest_text(
        text="Refunds take 5 days.",
        source_id="test_doc_1",
        product="RazorpayX",
        feature="General",
        access_level="INTERNAL"
    )
    
    collection = vector_db.get_collection()
    results = collection.get(where={"source_id": "test_doc_1"})
    
    doc_count = len(results["ids"])
    assert doc_count == 1, f"Expected 1 document, got {doc_count}"
    
    metadata = results["metadatas"][0]
    assert metadata["access_level"] == "INTERNAL", f"Expected INTERNAL, got {metadata['access_level']}"
    assert metadata["product"] == "RazorpayX", f"Expected RazorpayX, got {metadata['product']}"
    
    print(f"  [PASS] Document count: {doc_count}")
    print(f"  [PASS] access_level: {metadata['access_level']}")
    print(f"  [PASS] product: {metadata['product']}")
    print("  TEST CASE 1: PASSED")
    
    return True

def test_case_2_overwrite_logic():
    print("\n" + "="*60)
    print("TEST CASE 2: Overwrite Logic (Upsert)")
    print("="*60)
    
    chunk_count = vector_db.ingest_text(
        text="Refunds take 10 minutes.",
        source_id="test_doc_1",
        product="RazorpayX",
        feature="Payouts",
        access_level="EXTERNAL"
    )
    
    collection = vector_db.get_collection()
    results = collection.get(where={"source_id": "test_doc_1"})
    
    doc_count = len(results["ids"])
    assert doc_count == 1, f"Expected 1 document (not 2), got {doc_count}"
    
    content = results["documents"][0]
    assert "10 minutes" in content, f"Expected new content with '10 minutes', got: {content}"
    assert "5 days" not in content, f"Old content '5 days' should be deleted, but found in: {content}"
    
    print(f"  [PASS] Document count after overwrite: {doc_count} (not duplicated)")
    print(f"  [PASS] Content updated: '{content[:50]}...'")
    print("  TEST CASE 2: PASSED")
    
    return True

def test_case_3_delete_logic():
    print("\n" + "="*60)
    print("TEST CASE 3: Delete Logic")
    print("="*60)
    
    deleted_count = vector_db.delete_by_source_id("test_doc_1")
    print(f"  Deleted {deleted_count} chunks")
    
    collection = vector_db.get_collection()
    results = collection.get(where={"source_id": "test_doc_1"})
    
    doc_count = len(results["ids"])
    assert doc_count == 0, f"Expected 0 documents after deletion, got {doc_count}"
    
    print(f"  [PASS] Document count after deletion: {doc_count}")
    print("  TEST CASE 3: PASSED")
    
    return True

def test_case_4_access_level_filtering():
    print("\n" + "="*60)
    print("TEST CASE 4: Access Level Filtering (Security Check)")
    print("="*60)
    
    vector_db.ingest_text(
        text="Public Info - This is visible to everyone.",
        source_id="doc_public",
        product="RazorpayX",
        feature="General",
        access_level="EXTERNAL"
    )
    
    vector_db.ingest_text(
        text="Secret Info - This is private and confidential.",
        source_id="doc_private",
        product="RazorpayX",
        feature="General",
        access_level="INTERNAL"
    )
    
    collection = vector_db.get_collection()
    all_docs = collection.get()
    print(f"  Total documents in DB: {len(all_docs['ids'])}")
    
    results = collection.query(
        query_texts=["information"],
        where={"access_level": "EXTERNAL"},
        n_results=10
    )
    
    returned_docs = results["documents"][0] if results["documents"] else []
    returned_metadatas = results["metadatas"][0] if results["metadatas"] else []
    
    print(f"  Filtered query returned {len(returned_docs)} document(s)")
    
    for i, (doc, meta) in enumerate(zip(returned_docs, returned_metadatas)):
        assert meta["access_level"] == "EXTERNAL", f"Security breach! INTERNAL doc returned: {doc[:50]}"
        assert "Secret" not in doc, f"Security breach! Secret content visible: {doc[:50]}"
        print(f"    Doc {i+1}: access_level={meta['access_level']}, content='{doc[:30]}...'")
    
    has_public = any("Public" in doc for doc in returned_docs)
    assert has_public, "Public document not found in results"
    
    print("  [PASS] Only EXTERNAL documents returned")
    print("  [PASS] INTERNAL documents properly hidden")
    print("  TEST CASE 4: PASSED")
    
    return True

def run_all_tests():
    print("\n" + "#"*60)
    print("#  RAG CHATBOT - VECTOR DATABASE TEST SUITE")
    print("#"*60)
    
    all_passed = True
    
    try:
        if not test_case_1_fresh_ingestion():
            all_passed = False
        
        if not test_case_2_overwrite_logic():
            all_passed = False
        
        if not test_case_3_delete_logic():
            all_passed = False
        
        if not test_case_4_access_level_filtering():
            all_passed = False
        
    except AssertionError as e:
        print(f"\n  [FAIL] Assertion Error: {e}")
        all_passed = False
    except Exception as e:
        print(f"\n  [FAIL] Unexpected Error: {e}")
        all_passed = False
    finally:
        cleanup_test_db()
    
    print("\n" + "#"*60)
    if all_passed:
        print("#  ALL SYSTEMS GO: Database Logic is Secure")
        print("#"*60)
        return 0
    else:
        print("#  SOME TESTS FAILED - Please review the errors above")
        print("#"*60)
        return 1

if __name__ == "__main__":
    sys.exit(run_all_tests())