251 lines
8.9 KiB
Python
251 lines
8.9 KiB
Python
|
"""Tests for api_moneyflow module.
|
||
|
|
|
||
|
|
Tests the moneyflow (个股资金流向) API wrapper.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import pandas as pd
|
||
|
|
from unittest.mock import patch, MagicMock
|
||
|
|
|
||
|
|
from src.data.api_wrappers.api_moneyflow import (
|
||
|
|
get_moneyflow,
|
||
|
|
sync_moneyflow,
|
||
|
|
preview_moneyflow_sync,
|
||
|
|
MoneyflowSync,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TestMoneyflowAPI:
|
||
|
|
"""Test suite for moneyflow API wrapper."""
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_get_moneyflow_by_date(self, mock_client_class):
|
||
|
|
"""Test fetching moneyflow data by date."""
|
||
|
|
# Setup mock
|
||
|
|
mock_client = MagicMock()
|
||
|
|
mock_client_class.return_value = mock_client
|
||
|
|
mock_client.query.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ", "000002.SZ"],
|
||
|
|
"trade_date": ["20240115", "20240115"],
|
||
|
|
"buy_sm_vol": [10000, 20000],
|
||
|
|
"buy_sm_amount": [100.5, 200.5],
|
||
|
|
"sell_sm_vol": [8000, 15000],
|
||
|
|
"sell_sm_amount": [80.5, 150.5],
|
||
|
|
"buy_md_vol": [5000, 10000],
|
||
|
|
"buy_md_amount": [500.5, 1000.5],
|
||
|
|
"sell_md_vol": [4000, 8000],
|
||
|
|
"sell_md_amount": [400.5, 800.5],
|
||
|
|
"buy_lg_vol": [2000, 5000],
|
||
|
|
"buy_lg_amount": [2000.5, 5000.5],
|
||
|
|
"sell_lg_vol": [1500, 4000],
|
||
|
|
"sell_lg_amount": [1500.5, 4000.5],
|
||
|
|
"buy_elg_vol": [1000, 3000],
|
||
|
|
"buy_elg_amount": [5000.5, 15000.5],
|
||
|
|
"sell_elg_vol": [800, 2500],
|
||
|
|
"sell_elg_amount": [4000.5, 12500.5],
|
||
|
|
"net_mf_vol": [2700, 8000],
|
||
|
|
"net_mf_amount": [220.0, 550.0],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = get_moneyflow(trade_date="20240115")
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert not result.empty
|
||
|
|
assert len(result) == 2
|
||
|
|
assert "ts_code" in result.columns
|
||
|
|
assert "trade_date" in result.columns
|
||
|
|
assert "buy_sm_vol" in result.columns
|
||
|
|
assert "net_mf_amount" in result.columns
|
||
|
|
mock_client.query.assert_called_once_with("moneyflow", trade_date="20240115")
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_get_moneyflow_by_stock(self, mock_client_class):
|
||
|
|
"""Test fetching moneyflow data by stock code."""
|
||
|
|
# Setup mock
|
||
|
|
mock_client = MagicMock()
|
||
|
|
mock_client_class.return_value = mock_client
|
||
|
|
mock_client.query.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ", "000001.SZ"],
|
||
|
|
"trade_date": ["20240114", "20240115"],
|
||
|
|
"buy_sm_vol": [10000, 11000],
|
||
|
|
"buy_sm_amount": [100.5, 110.5],
|
||
|
|
"sell_sm_vol": [8000, 8500],
|
||
|
|
"sell_sm_amount": [80.5, 85.5],
|
||
|
|
"buy_md_vol": [5000, 5500],
|
||
|
|
"buy_md_amount": [500.5, 550.5],
|
||
|
|
"sell_md_vol": [4000, 4200],
|
||
|
|
"sell_md_amount": [400.5, 420.5],
|
||
|
|
"buy_lg_vol": [2000, 2200],
|
||
|
|
"buy_lg_amount": [2000.5, 2200.5],
|
||
|
|
"sell_lg_vol": [1500, 1600],
|
||
|
|
"sell_lg_amount": [1500.5, 1600.5],
|
||
|
|
"buy_elg_vol": [1000, 1100],
|
||
|
|
"buy_elg_amount": [5000.5, 5500.5],
|
||
|
|
"sell_elg_vol": [800, 850],
|
||
|
|
"sell_elg_amount": [4000.5, 4250.5],
|
||
|
|
"net_mf_vol": [2700, 2950],
|
||
|
|
"net_mf_amount": [220.0, 245.0],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = get_moneyflow(
|
||
|
|
ts_code="000001.SZ", start_date="20240101", end_date="20240131"
|
||
|
|
)
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert not result.empty
|
||
|
|
assert len(result) == 2
|
||
|
|
mock_client.query.assert_called_once_with(
|
||
|
|
"moneyflow", ts_code="000001.SZ", start_date="20240101", end_date="20240131"
|
||
|
|
)
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_get_moneyflow_empty_response(self, mock_client_class):
|
||
|
|
"""Test handling empty response."""
|
||
|
|
# Setup mock
|
||
|
|
mock_client = MagicMock()
|
||
|
|
mock_client_class.return_value = mock_client
|
||
|
|
mock_client.query.return_value = pd.DataFrame()
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = get_moneyflow(trade_date="20240101")
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert result.empty
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_get_moneyflow_with_shared_client(self, mock_client_class):
|
||
|
|
"""Test fetching with shared client for rate limiting."""
|
||
|
|
# Setup mock shared client
|
||
|
|
mock_shared_client = MagicMock()
|
||
|
|
mock_shared_client.query.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ"],
|
||
|
|
"trade_date": ["20240115"],
|
||
|
|
"buy_sm_vol": [10000],
|
||
|
|
"net_mf_amount": [220.0],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test with shared client
|
||
|
|
result = get_moneyflow(trade_date="20240115", client=mock_shared_client)
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert not result.empty
|
||
|
|
mock_shared_client.query.assert_called_once_with(
|
||
|
|
"moneyflow", trade_date="20240115"
|
||
|
|
)
|
||
|
|
mock_client_class.assert_not_called() # Should not create new client
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_get_moneyflow_date_column_rename(self, mock_client_class):
|
||
|
|
"""Test that 'date' column is renamed to 'trade_date'."""
|
||
|
|
# Setup mock with 'date' column (should be renamed)
|
||
|
|
mock_client = MagicMock()
|
||
|
|
mock_client_class.return_value = mock_client
|
||
|
|
mock_client.query.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ"],
|
||
|
|
"date": ["20240115"], # Using 'date' instead of 'trade_date'
|
||
|
|
"buy_sm_vol": [10000],
|
||
|
|
"net_mf_amount": [220.0],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = get_moneyflow(trade_date="20240115")
|
||
|
|
|
||
|
|
# Assert - date column should be renamed to trade_date
|
||
|
|
assert "trade_date" in result.columns
|
||
|
|
assert "date" not in result.columns
|
||
|
|
assert result["trade_date"].iloc[0] == "20240115"
|
||
|
|
|
||
|
|
|
||
|
|
class TestMoneyflowSync:
|
||
|
|
"""Test suite for MoneyflowSync class."""
|
||
|
|
|
||
|
|
def test_moneyflow_sync_class_attributes(self):
|
||
|
|
"""Test MoneyflowSync class has correct attributes."""
|
||
|
|
assert MoneyflowSync.table_name == "moneyflow"
|
||
|
|
assert MoneyflowSync.default_start_date == "20100101"
|
||
|
|
assert "ts_code" in MoneyflowSync.TABLE_SCHEMA
|
||
|
|
assert "trade_date" in MoneyflowSync.TABLE_SCHEMA
|
||
|
|
assert "buy_sm_vol" in MoneyflowSync.TABLE_SCHEMA
|
||
|
|
assert "net_mf_amount" in MoneyflowSync.TABLE_SCHEMA
|
||
|
|
assert MoneyflowSync.PRIMARY_KEY == ("ts_code", "trade_date")
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.get_moneyflow")
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.TushareClient")
|
||
|
|
def test_fetch_single_date(self, mock_client_class, mock_get_moneyflow):
|
||
|
|
"""Test fetch_single_date method."""
|
||
|
|
# Setup mock
|
||
|
|
mock_get_moneyflow.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ"],
|
||
|
|
"trade_date": ["20240115"],
|
||
|
|
"buy_sm_vol": [10000],
|
||
|
|
"net_mf_amount": [220.0],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Create sync instance
|
||
|
|
sync = MoneyflowSync()
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = sync.fetch_single_date("20240115")
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert not result.empty
|
||
|
|
mock_get_moneyflow.assert_called_once_with(
|
||
|
|
trade_date="20240115", client=sync.client
|
||
|
|
)
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.MoneyflowSync.sync_all")
|
||
|
|
def test_sync_moneyflow_function(self, mock_sync_all):
|
||
|
|
"""Test sync_moneyflow convenience function."""
|
||
|
|
# Setup mock
|
||
|
|
mock_sync_all.return_value = pd.DataFrame(
|
||
|
|
{
|
||
|
|
"ts_code": ["000001.SZ"],
|
||
|
|
"trade_date": ["20240115"],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = sync_moneyflow(start_date="20240101", end_date="20240131")
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert not result.empty
|
||
|
|
mock_sync_all.assert_called_once_with(
|
||
|
|
start_date="20240101", end_date="20240131", force_full=False
|
||
|
|
)
|
||
|
|
|
||
|
|
@patch("src.data.api_wrappers.api_moneyflow.MoneyflowSync.preview_sync")
|
||
|
|
def test_preview_moneyflow_sync_function(self, mock_preview_sync):
|
||
|
|
"""Test preview_moneyflow_sync convenience function."""
|
||
|
|
# Setup mock
|
||
|
|
mock_preview_sync.return_value = {
|
||
|
|
"sync_needed": True,
|
||
|
|
"date_count": 31,
|
||
|
|
"start_date": "20240101",
|
||
|
|
"end_date": "20240131",
|
||
|
|
"estimated_records": 10000,
|
||
|
|
"sample_data": pd.DataFrame(),
|
||
|
|
"mode": "incremental",
|
||
|
|
}
|
||
|
|
|
||
|
|
# Test
|
||
|
|
result = preview_moneyflow_sync(start_date="20240101", end_date="20240131")
|
||
|
|
|
||
|
|
# Assert
|
||
|
|
assert result["sync_needed"] is True
|
||
|
|
assert result["date_count"] == 31
|
||
|
|
mock_preview_sync.assert_called_once_with(
|
||
|
|
start_date="20240101", end_date="20240131", force_full=False, sample_size=3
|
||
|
|
)
|