ProStock/tests/test_daily_storage.py

"""Tests for data/daily.h5 storage validation.

Validates two key points:
1. All stocks from stock_basic.csv are saved in daily.h5
2. No abnormal data with very few data points (< 10 rows per stock)
"""

import pytest
import pandas as pd
from pathlib import Path
from src.data.storage import Storage
from src.data.stock_basic import _get_csv_path


class TestDailyStorageValidation:
    """Test daily.h5 storage integrity and completeness."""

    @pytest.fixture
    def storage(self):
        """Create storage instance."""
        return Storage()

    @pytest.fixture
    def stock_basic_df(self):
        """Load stock basic data from CSV."""
        csv_path = _get_csv_path()
        if not csv_path.exists():
            pytest.skip(f"stock_basic.csv not found at {csv_path}")
        return pd.read_csv(csv_path)

    @pytest.fixture
    def daily_df(self, storage):
        """Load daily data from HDF5."""
        if not storage.exists("daily"):
            pytest.skip("daily.h5 not found")
        # HDF5 stores keys with leading slash, so we need to handle both '/daily' and 'daily'
        file_path = storage._get_file_path("daily")
        try:
            with pd.HDFStore(file_path, mode="r") as store:
                if "/daily" in store.keys():
                    return store["/daily"]
                elif "daily" in store.keys():
                    return store["daily"]
                return pd.DataFrame()
        except Exception as e:
            pytest.skip(f"Error loading daily.h5: {e}")

    def test_all_stocks_saved(self, storage, stock_basic_df, daily_df):
        """Verify all stocks from stock_basic are saved in daily.h5.

        This test ensures data completeness - every stock in stock_basic
        should have corresponding data in daily.h5.
        """
        if daily_df.empty:
            pytest.fail("daily.h5 is empty")

        # Get unique stock codes from both sources
        expected_codes = set(stock_basic_df["ts_code"].dropna().unique())
        actual_codes = set(daily_df["ts_code"].dropna().unique())

        # Check for missing stocks
        missing_codes = expected_codes - actual_codes

        if missing_codes:
            missing_list = sorted(missing_codes)
            # Show first 20 missing stocks as sample
            sample = missing_list[:20]
            msg = f"Found {len(missing_codes)} stocks missing from daily.h5:\n"
            msg += f"Sample missing: {sample}\n"
            if len(missing_list) > 20:
                msg += f"... and {len(missing_list) - 20} more"
            pytest.fail(msg)

        # All stocks present
        assert len(actual_codes) > 0, "No stocks found in daily.h5"
        print(
            f"[TEST] All {len(expected_codes)} stocks from stock_basic are present in daily.h5"
        )

    def test_no_stock_with_insufficient_data(self, storage, daily_df):
        """Verify no stock has abnormally few data points (< 10 rows).

        Stocks with very few data points may indicate sync failures,
        delisted stocks not properly handled, or data corruption.
        """
        if daily_df.empty:
            pytest.fail("daily.h5 is empty")

        # Count rows per stock
        stock_counts = daily_df.groupby("ts_code").size()

        # Find stocks with less than 10 data points
        insufficient_stocks = stock_counts[stock_counts < 10]

        if not insufficient_stocks.empty:
            # Separate into categories for better reporting
            empty_stocks = stock_counts[stock_counts == 0]
            very_few_stocks = stock_counts[(stock_counts > 0) & (stock_counts < 10)]

            msg = f"Found {len(insufficient_stocks)} stocks with insufficient data (< 10 rows):\n"

            if not empty_stocks.empty:
                msg += f"\nEmpty stocks (0 rows): {len(empty_stocks)}\n"
                sample = sorted(empty_stocks.index[:10].tolist())
                msg += f"Sample: {sample}"

            if not very_few_stocks.empty:
                msg += f"\nVery few data points (1-9 rows): {len(very_few_stocks)}\n"
                # Show counts for these stocks
                sample = very_few_stocks.sort_values().head(20)
                msg += "Sample (ts_code: count):\n"
                for code, count in sample.items():
                    msg += f"  {code}: {count} rows\n"

            pytest.fail(msg)

        print(f"[TEST] All stocks have sufficient data (>= 10 rows)")

    def test_data_integrity_basic(self, storage, daily_df):
        """Basic data integrity checks for daily.h5."""
        if daily_df.empty:
            pytest.fail("daily.h5 is empty")

        # Check required columns exist
        required_columns = ["ts_code", "trade_date"]
        missing_columns = [
            col for col in required_columns if col not in daily_df.columns
        ]

        if missing_columns:
            pytest.fail(f"Missing required columns: {missing_columns}")

        # Check for null values in key columns
        null_ts_code = daily_df["ts_code"].isna().sum()
        null_trade_date = daily_df["trade_date"].isna().sum()

        if null_ts_code > 0:
            pytest.fail(f"Found {null_ts_code} rows with null ts_code")
        if null_trade_date > 0:
            pytest.fail(f"Found {null_trade_date} rows with null trade_date")

        print(f"[TEST] Data integrity check passed")

    def test_stock_data_coverage_report(self, storage, daily_df):
        """Generate a summary report of stock data coverage.

        This test provides visibility into data distribution without failing.
        """
        if daily_df.empty:
            pytest.skip("daily.h5 is empty - cannot generate report")

        stock_counts = daily_df.groupby("ts_code").size()

        # Calculate statistics
        total_stocks = len(stock_counts)
        min_count = stock_counts.min()
        max_count = stock_counts.max()
        median_count = stock_counts.median()
        mean_count = stock_counts.mean()

        # Distribution buckets
        very_low = (stock_counts < 10).sum()
        low = ((stock_counts >= 10) & (stock_counts < 100)).sum()
        medium = ((stock_counts >= 100) & (stock_counts < 500)).sum()
        high = (stock_counts >= 500).sum()

        report = f"""
=== Stock Data Coverage Report ===
Total stocks: {total_stocks}
Data points per stock:
  Min: {min_count}
  Max: {max_count}
  Median: {median_count:.0f}
  Mean: {mean_count:.1f}

Distribution:
  < 10 rows:  {very_low} stocks ({very_low / total_stocks * 100:.1f}%)
  10-99:      {low} stocks ({low / total_stocks * 100:.1f}%)
  100-499:    {medium} stocks ({medium / total_stocks * 100:.1f}%)
  >= 500:     {high} stocks ({high / total_stocks * 100:.1f}%)
"""
        print(report)

        # This is an informational test - it should not fail
        # But we assert to mark it as passed
        assert total_stocks > 0


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])