test_external_cache.py

#!/usr/bin/env python3
"""Test script to verify external SSD cache configuration works correctly."""

import os
import sys
from pathlib import Path


def test_cache_configuration():
    """Test that the external cache configuration is working."""

    print("🧪 Testing External SSD Cache Configuration")
    print("=" * 50)

    # Check if external SSD is mounted
    external_path = Path("/Volumes/extssd")
    if not external_path.exists():
        print("❌ External SSD not found at /Volumes/extssd")
        return False

    print("✅ External SSD is mounted")

    # Check if HF_HOME is set correctly
    hf_home = os.environ.get("HF_HOME")
    expected_hf_home = "/Volumes/extssd/huggingface"

    if hf_home != expected_hf_home:
        print(
            f"⚠️  HF_HOME not set correctly. Expected: {expected_hf_home}, Got: {hf_home}"
        )
        print("   Set HF_HOME with: export HF_HOME=/Volumes/extssd/huggingface")
        return False

    print(f"✅ HF_HOME correctly set to: {hf_home}")

    # Check if cache directories exist
    hub_cache = Path(hf_home) / "hub"
    if not hub_cache.exists():
        print(f"❌ Hub cache directory not found at: {hub_cache}")
        return False

    print(f"✅ Hub cache directory exists at: {hub_cache}")

    # Check if models are present
    model_count = len(list(hub_cache.glob("models--*")))
    print(f"✅ Found {model_count} models in cache")

    # Test importing Hugging Face libraries and check their cache detection
    try:
        from transformers import AutoTokenizer

        print("✅ Hugging Face libraries imported successfully")

        # Test a small model to verify cache is working
        print("🔄 Testing cache with a small model (this may take a moment)...")

        # This should use the external cache
        _ = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")

        print("✅ Successfully loaded model from cache")

        # Check if the model files are in the expected location
        clip_path = hub_cache / "models--openai--clip-vit-base-patch32"
        if clip_path.exists():
            print(f"✅ Model files found in external cache at: {clip_path}")
        else:
            print(f"⚠️  Model files not found at expected location: {clip_path}")
            return False

        return True

    except Exception as e:
        print(f"❌ Error loading model: {e}")
        return False


def main():
    """Main test function."""
    # Load .env file if available
    env_file = Path(__file__).parent / ".env"
    if env_file.exists():
        print("📁 Loading .env file...")
        from dotenv import load_dotenv

        load_dotenv()
    else:
        print("⚠️  No .env file found, using system environment variables")

    success = test_cache_configuration()

    print("\n" + "=" * 50)
    if success:
        print("🎉 All tests passed! External SSD cache is working correctly.")
        print("You can now run the application with: ./run.sh")
    else:
        print("❌ Some tests failed. Please check the configuration.")
        sys.exit(1)


if __name__ == "__main__":
    main()