CVE-FactChecker / test_production_simulation.py
NLPGenius's picture
Fix deployment issues: enhanced environment config, robust background ingestion, improved health checks, production-ready
aa69d4c
Raw
History Blame Contribute Delete
12.5 kB
#!/usr/bin/env python3
"""
Production simulation test to validate all deployment fixes.
Simulates the HuggingFace Spaces environment and tests the complete system.
"""
import os
import sys
import time
import subprocess
import threading
import requests
from datetime import datetime
# Add the parent directory to Python path
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)
class ProductionSimulator:
def __init__(self):
self.original_env = {}
self.test_results = {}
def setup_production_environment(self):
"""Set up environment variables to simulate HuggingFace Spaces deployment."""
print("πŸ”§ Setting up production environment simulation...")
# Store original environment
env_vars_to_set = {
"AUTO_INGEST": "true",
"LANGUAGE_FILTER": "English",
"PORT": "7860",
"HF_HOME": "/tmp/huggingface",
"TRANSFORMERS_CACHE": "/tmp/transformers",
"VECTOR_PERSIST_DIR": "/tmp/vector_db"
}
for key, value in env_vars_to_set.items():
self.original_env[key] = os.environ.get(key)
os.environ[key] = value
print(f" βœ… {key}={value}")
def cleanup_environment(self):
"""Restore original environment."""
print("🧹 Cleaning up environment...")
for key, original_value in self.original_env.items():
if original_value is None:
os.environ.pop(key, None)
else:
os.environ[key] = original_value
def test_app_initialization(self):
"""Test Flask app initialization with production settings."""
print("\nπŸ” Testing App Initialization...")
try:
# Clear any existing lock files
lock_file = "/tmp/ingest.lock" if os.name != 'nt' else "ingest.lock"
if os.path.exists(lock_file):
os.remove(lock_file)
print(" πŸ—‘οΈ Cleared existing lock file")
# Import and test app components
from cve_factchecker.app import app, _safe_initialize_system, INGEST_STATUS, AUTO_INGEST
print(f" πŸ“Š AUTO_INGEST: {AUTO_INGEST}")
print(f" πŸ“Š INGEST_STATUS: {INGEST_STATUS}")
# Test system initialization
_safe_initialize_system()
print(" βœ… System initialization completed")
# Check if background thread should start
from cve_factchecker.app import should_start_ingestion
print(f" πŸ“Š Should start ingestion: {should_start_ingestion}")
self.test_results["app_initialization"] = {
"success": True,
"auto_ingest_enabled": AUTO_INGEST,
"ingestion_status": INGEST_STATUS.copy()
}
except Exception as e:
print(f" ❌ App initialization failed: {e}")
self.test_results["app_initialization"] = {
"success": False,
"error": str(e)
}
return False
return True
def test_health_endpoint_behavior(self):
"""Test the enhanced health endpoint."""
print("\nπŸ” Testing Health Endpoint...")
try:
from cve_factchecker.app import app
with app.test_client() as client:
# Test basic health check
response = client.get('/health')
health_data = response.get_json()
print(f" πŸ“Š Health Status: {health_data.get('status')}")
print(f" πŸ“Š Vector Store Populated: {health_data.get('vector_store_populated', 'unknown')}")
print(f" πŸ“Š Sample Documents: {health_data.get('sample_documents', 0)}")
# Test ingestion trigger if vector store is empty
if not health_data.get('vector_store_populated', False):
print(" πŸ”„ Testing ingestion trigger...")
trigger_response = client.get('/health?trigger_ingestion=true')
trigger_data = trigger_response.get_json()
print(f" πŸ“Š Trigger Response: {trigger_data.get('message', 'No message')}")
self.test_results["health_endpoint"] = {
"success": True,
"health_data": health_data,
"trigger_tested": not health_data.get('vector_store_populated', False)
}
except Exception as e:
print(f" ❌ Health endpoint test failed: {e}")
self.test_results["health_endpoint"] = {
"success": False,
"error": str(e)
}
return False
return True
def test_background_ingestion_flow(self):
"""Test the complete background ingestion flow."""
print("\nπŸ” Testing Background Ingestion Flow...")
try:
from cve_factchecker.app import _background_ingest, INGEST_STATUS, _cleanup_stale_locks
# Test stale lock cleanup
print(" 🧹 Testing stale lock cleanup...")
_cleanup_stale_locks()
# Reset ingestion status
INGEST_STATUS.update({"finished": False, "test_mode": True})
# Run background ingestion in test mode
print(" πŸš€ Running background ingestion...")
start_time = time.time()
# Use a thread to avoid blocking
ingestion_thread = threading.Thread(target=_background_ingest, daemon=True)
ingestion_thread.start()
# Wait for completion with timeout
timeout = 60 # 1 minute timeout
while not INGEST_STATUS.get("finished") and (time.time() - start_time) < timeout:
time.sleep(1)
print(f" ⏳ Waiting for ingestion... ({time.time() - start_time:.0f}s)")
ingestion_time = time.time() - start_time
if INGEST_STATUS.get("finished"):
print(f" βœ… Ingestion completed in {ingestion_time:.1f}s")
print(f" πŸ“Š Synced articles: {INGEST_STATUS.get('synced', 0)}")
if INGEST_STATUS.get("error"):
print(f" ⚠️ Ingestion error: {INGEST_STATUS.get('error')}")
self.test_results["background_ingestion"] = {
"success": True,
"completion_time": ingestion_time,
"final_status": INGEST_STATUS.copy()
}
else:
print(f" ❌ Ingestion timed out after {timeout}s")
self.test_results["background_ingestion"] = {
"success": False,
"error": "Timeout",
"partial_status": INGEST_STATUS.copy()
}
return False
except Exception as e:
print(f" ❌ Background ingestion test failed: {e}")
self.test_results["background_ingestion"] = {
"success": False,
"error": str(e)
}
return False
return True
def test_fact_checking_after_ingestion(self):
"""Test fact-checking functionality after ingestion."""
print("\nπŸ” Testing Fact-Checking After Ingestion...")
try:
from cve_factchecker.app import app
with app.test_client() as client:
test_claims = [
"Security researchers discovered a new vulnerability",
"Cyberattack hits major corporation",
"Malware targets government systems"
]
for claim in test_claims:
print(f" πŸ” Testing claim: {claim[:50]}...")
response = client.post('/fact-check', json={"claim": claim})
result = response.get_json()
print(f" πŸ“Š Verdict: {result.get('verdict', 'Unknown')}")
print(f" πŸ“Š Sources: {result.get('sources_used', 0)}")
print(f" πŸ“Š Confidence: {result.get('confidence', 0)}")
if result.get('verdict') not in ['ERROR', 'INITIALIZING']:
print(f" βœ… Fact-check working")
break
else:
print(f" ❌ No successful fact-checks")
return False
self.test_results["fact_checking"] = {
"success": True,
"test_claims": len(test_claims),
"sample_result": result
}
except Exception as e:
print(f" ❌ Fact-checking test failed: {e}")
self.test_results["fact_checking"] = {
"success": False,
"error": str(e)
}
return False
return True
def test_production_simulation(self):
"""Run complete production simulation test."""
print("🏭 CVE Fact Checker - Production Simulation Test")
print("=" * 80)
success = True
try:
self.setup_production_environment()
# Run tests in sequence
tests = [
("App Initialization", self.test_app_initialization),
("Health Endpoint", self.test_health_endpoint_behavior),
("Background Ingestion", self.test_background_ingestion_flow),
("Fact-Checking", self.test_fact_checking_after_ingestion)
]
for test_name, test_func in tests:
print(f"\n{'='*20} {test_name} {'='*20}")
test_success = test_func()
success = success and test_success
if not test_success:
print(f"❌ {test_name} failed - stopping tests")
break
else:
print(f"βœ… {test_name} passed")
finally:
self.cleanup_environment()
return success
def print_summary(self):
"""Print test summary."""
print("\nπŸ“‹ Production Simulation Summary")
print("=" * 50)
total_tests = len(self.test_results)
passed_tests = sum(1 for result in self.test_results.values() if result.get("success"))
print(f"Tests Run: {total_tests}")
print(f"Tests Passed: {passed_tests}")
print(f"Tests Failed: {total_tests - passed_tests}")
for test_name, result in self.test_results.items():
status = "βœ… PASS" if result.get("success") else "❌ FAIL"
print(f"{status} {test_name.replace('_', ' ').title()}")
if not result.get("success") and result.get("error"):
print(f" Error: {result['error']}")
overall_success = passed_tests == total_tests
print(f"\nOverall Result: {'βœ… SUCCESS' if overall_success else '❌ FAILURE'}")
return overall_success
def main():
"""Main test function."""
simulator = ProductionSimulator()
try:
success = simulator.test_production_simulation()
simulator.print_summary()
if success:
print("\nπŸŽ‰ Production simulation successful!")
print("πŸ’‘ System is ready for deployment to HuggingFace Spaces")
else:
print("\n🚨 Production simulation failed!")
print("πŸ’‘ Issues need to be resolved before deployment")
return success
except KeyboardInterrupt:
print("\n⏹️ Test interrupted by user")
return False
except Exception as e:
print(f"\n❌ Test suite failed: {e}")
return False
if __name__ == "__main__":
success = main()
sys.exit(0 if success else 1)