sanguo_vnpy_v2/tests/data/test_index_downloader.py

"""Tests for index downloader and read_index_daily functionality."""
import pandas as pd
import os
from unittest.mock import patch, MagicMock
from datetime import date
from pathlib import Path
import pytest

from sanguo_data.config import DataConfig


def test_download_index_writes_parquet(tmp_path):
    """Test that download_index writes parquet files with correct structure."""
    # Sample data that baostock would return
    sample_data = [
        ["2024-01-02", "3495.0", "3505.0", "3490.0", "3500.0", "100000"],
        ["2024-01-03", "3505.0", "3515.0", "3500.0", "3510.0", "120000"],
        ["2024-01-04", "3515.0", "3525.0", "3510.0", "3520.0", "110000"],
    ]

    # Create a simple baostock mock
    class MockBaostock:
        class MockResult:
            def __init__(self, data):
                self.error_code = "success"
                self.error_msg = "success"
                self.data = data
                self.fields = ["date", "open", "high", "low", "close", "volume"]
                self.row_index = 0

            def next(self):
                if self.row_index < len(self.data):
                    row = self.data[self.row_index]
                    self.row_index += 1
                    return True
                return False

            def get_row_data(self):
                return self.data[self.row_index - 1]

        def login(self):
            return self.MockResult([])

        def logout(self):
            return self.MockResult([])

        def query_history_k_data_plus(self, *args, **kwargs):
            return self.MockResult(sample_data)

    # Patch baostock module
    import sys
    sys.modules["baostock"] = MockBaostock()

    try:
        # Import after patching
        from sanguo_data.index_downloader import download_index

        # Download index data
        download_index("sh000300", 2024, 2024, str(tmp_path))

    finally:
        # Clean up the mock
        del sys.modules["baostock"]

    # Verify parquet file was created
    expected_file = tmp_path / "2024" / "sh000300_daily.parquet"
    assert expected_file.exists(), f"Expected parquet file at {expected_file}"

    # Verify parquet content
    df_read = pd.read_parquet(expected_file)
    assert len(df_read) == 3
    assert "close" in df_read.columns
    assert "date" in df_read.columns
    assert df_read["close"].iloc[0] == 3500.0


def test_download_index_clears_proxy(tmp_path):
    """Test that download_index clears proxy environment variables."""
    # Set proxy variables
    os.environ["http_proxy"] = "http://evil:8080"
    os.environ["https_proxy"] = "https://evil:8080"

    sample_data = [["2024-01-02", "3495.0", "3505.0", "3490.0", "3500.0", "100000"]]

    # Create a simple baostock mock
    class MockBaostock:
        class MockResult:
            def __init__(self, data):
                self.error_code = "success"
                self.error_msg = "success"
                self.data = data
                self.fields = ["date", "open", "high", "low", "close", "volume"]
                self.row_index = 0

            def next(self):
                if self.row_index < len(self.data):
                    row = self.data[self.row_index]
                    self.row_index += 1
                    return True
                return False

            def get_row_data(self):
                return self.data[self.row_index - 1]

        def login(self):
            return self.MockResult([])

        def logout(self):
            return self.MockResult([])

        def query_history_k_data_plus(self, *args, **kwargs):
            return self.MockResult(sample_data)

    # Patch baostock module
    import sys
    sys.modules["baostock"] = MockBaostock()

    try:
        # Import after patching
        from sanguo_data.index_downloader import download_index

        # Download index data
        download_index("sh000300", 2024, 2024, str(tmp_path))

    finally:
        # Clean up the mock
        del sys.modules["baostock"]

    # Verify proxy variables were cleared
    assert "http_proxy" not in os.environ
    assert "https_proxy" not in os.environ


def test_read_index_daily_reads_parquet(tmp_path):
    """Test that read_index_daily reads index parquet files correctly."""
    # Import after implementation
    from sanguo_data.datareader import read_index_daily

    # Create test parquet file with same structure as stock data
    year_dir = tmp_path / "2024"
    year_dir.mkdir()

    df = pd.DataFrame({
        "date": ["2024-01-02", "2024-01-03", "2024-01-04"],
        "open": [3495.0, 3505.0, 3515.0],
        "high": [3505.0, 3515.0, 3525.0],
        "low": [3490.0, 3500.0, 3510.0],
        "close": [3500.0, 3510.0, 3520.0],
        "volume": [100000, 120000, 110000],
    })
    df.to_parquet(year_dir / "sh000300_daily.parquet")

    cfg = DataConfig(
        data_paths={"daily_dir": str(tmp_path)},
        data_sources={}, validation={}, performance={},
    )

    # Read index daily data
    result = read_index_daily("sh000300", date(2024, 1, 1), date(2024, 12, 31), cfg)

    # Verify result
    assert len(result) == 3
    assert "close" in result.columns
    assert result["close"].iloc[0] == 3500.0


def test_read_index_daily_handles_date_range(tmp_path):
    """Test that read_index_daily filters by date range correctly."""
    # Import after implementation
    from sanguo_data.datareader import read_index_daily

    # Create test parquet file
    year_dir = tmp_path / "2024"
    year_dir.mkdir()

    df = pd.DataFrame({
        "date": ["2024-01-02", "2024-06-15", "2024-12-31"],
        "open": [3495.0, 3600.0, 3700.0],
        "high": [3505.0, 3610.0, 3710.0],
        "low": [3490.0, 3590.0, 3690.0],
        "close": [3500.0, 3605.0, 3705.0],
        "volume": [100000, 120000, 110000],
    })
    df.to_parquet(year_dir / "sh000300_daily.parquet")

    cfg = DataConfig(
        data_paths={"daily_dir": str(tmp_path)},
        data_sources={}, validation={}, performance={},
    )

    # Read with narrowed date range
    result = read_index_daily("sh000300", date(2024, 1, 1), date(2024, 6, 30), cfg)

    # Should only get first 2 rows
    assert len(result) == 2
    assert result["close"].iloc[0] == 3500.0
    assert result["close"].iloc[1] == 3605.0