test_download.py

#!/usr/bin/env python3
"""
Test script to verify the data download functionality works properly.
This script removes existing data files and then attempts to download fresh ones.
"""
import os
import shutil
import sys
from download_from_hub import ensure_dirs, download_datasets, create_fallback_data

def clean_existing_data():
    """Remove existing data files to test a fresh download."""
    print("Cleaning existing data files...")
    
    # Paths to clean
    paths_to_remove = [
        "embeddings/faiss_index.index",
        "embeddings/embeddings.pkl",
        "data/doc_chunks.pkl"
    ]
    
    # Remove each file if it exists
    for path in paths_to_remove:
        if os.path.exists(path):
            print(f"Removing {path}")
            os.remove(path)
    
    print("Data files cleaned, ready for fresh download")

def test_download():
    """Test the complete download process."""
    # First ensure directories exist
    ensure_dirs()
    
    # Clean existing data
    clean_existing_data()
    
    # Try to download datasets
    print("\n===== STARTING DOWNLOAD TEST =====\n")
    success = download_datasets()
    
    if not success:
        print("\n===== DOWNLOAD FAILED, TESTING FALLBACK =====\n")
        success = create_fallback_data()
        if success:
            print("\n===== FALLBACK DATA CREATED SUCCESSFULLY =====\n")
        else:
            print("\n===== FALLBACK DATA CREATION FAILED =====\n")
            return False
    else:
        print("\n===== DOWNLOAD COMPLETED SUCCESSFULLY =====\n")
    
    # Verify all required files exist
    required_files = [
        "embeddings/faiss_index.index",
        "embeddings/embeddings.pkl",
        "data/doc_chunks.pkl"
    ]
    
    all_exist = True
    print("Verifying all required files exist:")
    for file_path in required_files:
        if os.path.exists(file_path):
            size_mb = os.path.getsize(file_path) / (1024 * 1024)
            print(f"✅ {file_path} exists (size: {size_mb:.2f} MB)")
        else:
            print(f"❌ {file_path} is missing")
            all_exist = False
    
    return all_exist

if __name__ == "__main__":
    if test_download():
        print("\n===== TEST PASSED: All data files were downloaded or created successfully =====")
        sys.exit(0)
    else:
        print("\n===== TEST FAILED: Some data files are missing =====")
        sys.exit(1)