pypi-query-mcp/test_improved.py

#!/usr/bin/env python3
"""Test script for the improved get_top_packages_by_downloads function."""

import asyncio

from pypi_query_mcp.tools.download_stats import get_top_packages_by_downloads


async def test_improved():
    try:
        result = await get_top_packages_by_downloads("month", 10)
        print("✅ Success! Result keys:", list(result.keys()))
        print(f"Number of packages returned: {len(result.get('top_packages', []))}")
        print(f"Data source: {result.get('data_source')}")
        print(f"Methodology: {result.get('methodology')}")

        print("\nTop 5 packages:")
        for i, pkg in enumerate(result.get("top_packages", [])[:5]):
            downloads = pkg.get("downloads", 0)
            stars = pkg.get("github_stars", "N/A")
            estimated = "(estimated)" if pkg.get("estimated", False) else "(real)"
            github_enhanced = " 🌟" if pkg.get("github_enhanced", False) else ""
            print(
                f"{i + 1}. {pkg.get('package', 'N/A')} - {downloads:,} downloads {estimated}{github_enhanced}"
            )
            if stars != "N/A":
                print(
                    f"   GitHub: {stars:,} stars, {pkg.get('category', 'N/A')} category"
                )

        # Test different periods
        print("\n--- Testing different periods ---")
        for period in ["day", "week", "month"]:
            result = await get_top_packages_by_downloads(period, 3)
            top_3 = result.get("top_packages", [])
            print(
                f"{period}: {len(top_3)} packages, avg downloads: {sum(p.get('downloads', 0) for p in top_3) // max(len(top_3), 1):,}"
            )

        print("\n--- Testing different limits ---")
        for limit in [5, 20, 50]:
            result = await get_top_packages_by_downloads("month", limit)
            packages = result.get("top_packages", [])
            real_count = len([p for p in packages if not p.get("estimated", False)])
            print(
                f"Limit {limit}: {len(packages)} packages returned, {real_count} with real stats"
            )

    except Exception as e:
        print(f"❌ Error: {e}")
        import traceback

        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(test_improved())