lib-x17-fintech 2.1.3__py3-none-any.whl

xfintech/data/source/tushare/stock/companyprofit/tests/test_class_companyprofit_all.py

@@ -0,0 +1,635 @@
+from datetime import datetime
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+import pytest
+
+from xfintech.data.common.cache import Cache
+from xfintech.data.common.params import Params
+from xfintech.data.source.tushare.session.session import Session
+from xfintech.data.source.tushare.stock.companyprofit.companyprofit import CompanyProfit
+from xfintech.data.source.tushare.stock.companyprofit.constant import (
+    KEY,
+    NAME,
+    PAGINATE,
+    SOURCE,
+    TARGET,
+)
+
+# ============================================================================
+# Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def mock_session():
+    """Create a mock Tushare session"""
+    session = MagicMock(spec=Session)
+    session._credential = "test_token"
+    session.id = "test1234"
+    session.mode = "direct"
+    session.relay_url = None
+    session.relay_secret = None
+    session.connected = True
+
+    # Mock the connection object
+    mock_connection = MagicMock()
+    mock_connection.income_vip = MagicMock()
+    session.connection = mock_connection
+
+    return session
+
+
+@pytest.fixture
+def sample_source_data():
+    """Create sample source data in Tushare format"""
+    return pd.DataFrame(
+        {
+            "ts_code": ["000001.SZ", "000001.SZ", "000001.SZ"],
+            "ann_date": ["20200425", "20200826", "20201030"],
+            "f_ann_date": ["20200425", "20200826", "20201030"],
+            "end_date": ["20200331", "20200630", "20200930"],
+            "report_type": ["1", "1", "1"],
+            "comp_type": ["1", "1", "1"],
+            "end_type": ["12", "12", "12"],
+            "update_flag": ["0", "0", "0"],
+            "basic_eps": [0.5, 1.2, 1.8],
+            "diluted_eps": [0.5, 1.2, 1.8],
+            "total_revenue": [100000000.0, 250000000.0, 380000000.0],
+            "revenue": [98000000.0, 245000000.0, 372000000.0],
+            "n_income": [12000000.0, 32000000.0, 48000000.0],
+            "n_income_attr_p": [11500000.0, 31000000.0, 46000000.0],
+            "total_profit": [15000000.0, 39000000.0, 58000000.0],
+            "income_tax": [3000000.0, 7000000.0, 10000000.0],
+        }
+    )
+
+
+@pytest.fixture
+def expected_transformed_data():
+    """Create expected transformed data"""
+    return pd.DataFrame(
+        {
+            "code": ["000001.SZ", "000001.SZ", "000001.SZ"],
+            "date": ["2020-03-31", "2020-06-30", "2020-09-30"],
+            "datecode": ["20200331", "20200630", "20200930"],
+            "ann_date": ["2020-04-25", "2020-08-26", "2020-10-30"],
+            "f_ann_date": ["2020-04-25", "2020-08-26", "2020-10-30"],
+            "end_date": ["2020-03-31", "2020-06-30", "2020-09-30"],
+            "report_type": ["1", "1", "1"],
+            "comp_type": ["1", "1", "1"],
+            "end_type": ["12", "12", "12"],
+            "update_flag": ["0", "0", "0"],
+            "basic_eps": [0.5, 1.2, 1.8],
+            "diluted_eps": [0.5, 1.2, 1.8],
+            "total_revenue": [100000000.0, 250000000.0, 380000000.0],
+            "revenue": [98000000.0, 245000000.0, 372000000.0],
+            "n_income": [12000000.0, 32000000.0, 48000000.0],
+            "n_income_attr_p": [11500000.0, 31000000.0, 46000000.0],
+            "total_profit": [15000000.0, 39000000.0, 58000000.0],
+            "income_tax": [3000000.0, 7000000.0, 10000000.0],
+        }
+    )
+
+
+# ============================================================================
+# Initialization Tests
+# ============================================================================
+
+
+def test_companyprofit_init_basic(mock_session):
+    """Test basic initialization with required session"""
+    profit = CompanyProfit(session=mock_session)
+    assert profit.name == NAME
+    assert profit.key == KEY
+    assert profit.source == SOURCE
+    assert profit.target == TARGET
+    assert profit.paginate == PAGINATE
+
+
+def test_companyprofit_init_with_params_dict(mock_session):
+    """Test initialization with params as dict"""
+    params = {"ts_code": "000001.SZ", "start_date": "20200101", "end_date": "20201231"}
+    profit = CompanyProfit(session=mock_session, params=params)
+    assert profit.params.get("ts_code") == "000001.SZ"
+    assert profit.params.get("start_date") == "20200101"
+    assert profit.params.get("end_date") == "20201231"
+
+
+def test_companyprofit_init_with_params_object(mock_session):
+    """Test initialization with params as Params object"""
+    params = Params(ts_code="000001.SZ", start_date="20200101")
+    profit = CompanyProfit(session=mock_session, params=params)
+    assert profit.params.get("ts_code") == "000001.SZ"
+
+
+def test_companyprofit_init_with_year_param(mock_session):
+    """Test initialization with year parameter"""
+    params = {"ts_code": "000001.SZ", "year": "2020"}
+    profit = CompanyProfit(session=mock_session, params=params)
+    assert profit.params.get("year") == "2020"
+
+
+def test_companyprofit_init_with_cache_bool_true(mock_session):
+    """Test initialization with cache as boolean True"""
+    profit = CompanyProfit(session=mock_session, cache=True)
+    assert profit.cache is not None
+    assert isinstance(profit.cache, Cache)
+
+
+def test_companyprofit_init_with_cache_bool_false(mock_session):
+    """Test initialization with cache as boolean False"""
+    profit = CompanyProfit(session=mock_session, cache=False)
+    assert profit.cache is None
+
+
+def test_companyprofit_init_with_cache_dict(mock_session):
+    """Test initialization with cache as dict"""
+    cache_config = {"dir": "/tmp/test_cache", "ttl": 3600}
+    profit = CompanyProfit(session=mock_session, cache=cache_config)
+    assert profit.cache is not None
+    assert isinstance(profit.cache, Cache)
+
+
+def test_companyprofit_init_default_paginate_limit(mock_session):
+    """Test that default paginate settings are correctly set"""
+    profit = CompanyProfit(session=mock_session)
+    assert profit.paginate.pagesize == 1000
+    assert profit.paginate.pagelimit == 1000
+
+
+def test_companyprofit_init_with_all_params(mock_session):
+    """Test initialization with all parameters"""
+    params = {"ts_code": "000001.SZ", "start_date": "20200101", "end_date": "20201231"}
+    coolant = {"interval": 0.5}
+    retry = {"max_attempts": 3}
+    cache = True
+
+    profit = CompanyProfit(session=mock_session, params=params, coolant=coolant, retry=retry, cache=cache)
+    assert profit.params.get("ts_code") == "000001.SZ"
+    assert profit.cache is not None
+
+
+# ============================================================================
+# Transform Tests
+# ============================================================================
+
+
+def test_companyprofit_transform_basic(mock_session, sample_source_data):
+    """Test basic transformation of source data"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    assert not result.empty
+    assert len(result) == 3
+    assert "code" in result.columns
+    assert "date" in result.columns
+    assert "datecode" in result.columns
+
+
+def test_companyprofit_transform_code_mapping(mock_session, sample_source_data):
+    """Test that ts_code is correctly mapped to code"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    assert all(result["code"] == ["000001.SZ", "000001.SZ", "000001.SZ"])
+
+
+def test_companyprofit_transform_date_format(mock_session, sample_source_data):
+    """Test that dates are formatted correctly (YYYY-MM-DD)"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    # Check date format
+    assert result["date"].iloc[0] == "2020-03-31"
+    assert result["ann_date"].iloc[0] == "2020-04-25"
+    assert result["f_ann_date"].iloc[0] == "2020-04-25"
+    assert result["end_date"].iloc[0] == "2020-03-31"
+
+
+def test_companyprofit_transform_datecode_preserved(mock_session, sample_source_data):
+    """Test that datecode is preserved in YYYYMMDD format"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    assert result["datecode"].iloc[0] == "20200331"
+    assert result["datecode"].iloc[1] == "20200630"
+    assert result["datecode"].iloc[2] == "20200930"
+
+
+def test_companyprofit_transform_numeric_conversions(mock_session, sample_source_data):
+    """Test that numeric fields are properly converted"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    # Check numeric fields
+    assert result["basic_eps"].dtype == "float64"
+    assert result["diluted_eps"].dtype == "float64"
+    assert result["total_revenue"].dtype == "float64"
+    assert result["n_income"].dtype == "float64"
+
+
+def test_companyprofit_transform_string_fields(mock_session, sample_source_data):
+    """Test that string fields are properly converted"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    assert result["report_type"].dtype == "object"
+    assert result["comp_type"].dtype == "object"
+    assert result["end_type"].dtype == "object"
+    assert result["update_flag"].dtype == "object"
+
+
+def test_companyprofit_transform_empty_dataframe(mock_session):
+    """Test transform with empty DataFrame"""
+    profit = CompanyProfit(session=mock_session)
+    empty_df = pd.DataFrame()
+    result = profit.transform(empty_df)
+
+    assert result.empty
+    assert len(result.columns) == len(TARGET.list_column_names())
+
+
+def test_companyprofit_transform_none_input(mock_session):
+    """Test transform with None input"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(None)
+
+    assert result.empty
+    assert len(result.columns) == len(TARGET.list_column_names())
+
+
+def test_companyprofit_transform_handles_invalid_dates(mock_session):
+    """Test transform handles invalid dates gracefully"""
+    profit = CompanyProfit(session=mock_session)
+    data = pd.DataFrame(
+        {
+            "ts_code": ["000001.SZ"],
+            "end_date": ["invalid"],
+            "ann_date": ["20200101"],
+            "f_ann_date": ["20200101"],
+            "report_type": ["1"],
+            "comp_type": ["1"],
+            "end_type": ["12"],
+            "update_flag": ["0"],
+        }
+    )
+    result = profit.transform(data)
+
+    assert not result.empty
+    assert pd.isna(result["date"].iloc[0]) or result["date"].iloc[0] == "NaT"
+
+
+def test_companyprofit_transform_removes_duplicates(mock_session):
+    """Test that duplicates are removed"""
+    profit = CompanyProfit(session=mock_session)
+    data = pd.DataFrame(
+        {
+            "ts_code": ["000001.SZ", "000001.SZ"],
+            "end_date": ["20200331", "20200331"],
+            "ann_date": ["20200425", "20200425"],
+            "f_ann_date": ["20200425", "20200425"],
+            "report_type": ["1", "1"],
+            "comp_type": ["1", "1"],
+            "end_type": ["12", "12"],
+            "update_flag": ["0", "0"],
+            "basic_eps": [0.5, 0.5],
+        }
+    )
+    result = profit.transform(data)
+
+    assert len(result) == 1
+
+
+def test_companyprofit_transform_sorts_by_code_and_date(mock_session):
+    """Test that output is sorted by code and date"""
+    profit = CompanyProfit(session=mock_session)
+    data = pd.DataFrame(
+        {
+            "ts_code": ["000002.SZ", "000001.SZ", "000001.SZ"],
+            "end_date": ["20200331", "20200630", "20200331"],
+            "ann_date": ["20200425", "20200826", "20200425"],
+            "f_ann_date": ["20200425", "20200826", "20200425"],
+            "report_type": ["1", "1", "1"],
+            "comp_type": ["1", "1", "1"],
+            "end_type": ["12", "12", "12"],
+            "update_flag": ["0", "0", "0"],
+        }
+    )
+    result = profit.transform(data)
+
+    assert result["code"].iloc[0] == "000001.SZ"
+    assert result["code"].iloc[1] == "000001.SZ"
+    assert result["code"].iloc[2] == "000002.SZ"
+    assert result["date"].iloc[0] == "2020-03-31"
+    assert result["date"].iloc[1] == "2020-06-30"
+
+
+def test_companyprofit_transform_resets_index(mock_session, sample_source_data):
+    """Test that index is reset after transformation"""
+    profit = CompanyProfit(session=mock_session)
+    result = profit.transform(sample_source_data)
+
+    assert list(result.index) == [0, 1, 2]
+
+
+def test_companyprofit_transform_only_target_columns(mock_session, sample_source_data):
+    """Test that only target columns are returned"""
+    profit = CompanyProfit(session=mock_session)
+
+    # Add extra column to source data
+    data_with_extra = sample_source_data.copy()
+    data_with_extra["extra_column"] = "test"
+
+    result = profit.transform(data_with_extra)
+
+    assert "extra_column" not in result.columns
+    assert "ts_code" not in result.columns  # Should be renamed to "code"
+    assert "code" in result.columns
+
+
+# ============================================================================
+# Run Tests
+# ============================================================================
+
+
+def test_companyprofit_run_with_cache_hit(mock_session, sample_source_data):
+    """Test run with cache hit (no API call)"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ"},
+        cache=True,
+    )
+
+    # Mock cache load to return data
+    with patch.object(profit, "_load_cache", return_value=sample_source_data):
+        result = profit.run()
+
+        assert not result.empty
+        mock_session.connection.income_vip.assert_not_called()
+
+
+def test_companyprofit_run_basic_date_range(mock_session, sample_source_data):
+    """Test run with start_date and end_date"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ", "start_date": "20200101", "end_date": "20201231"},
+    )
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data):
+        with patch.object(profit, "_save_cache"):
+            result = profit.run()
+
+            assert not result.empty
+            assert "code" in result.columns
+
+
+def test_companyprofit_run_with_year_param(mock_session, sample_source_data):
+    """Test that year parameter is converted to start_date and end_date"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ", "year": "2020"})
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_save_cache"):
+            profit.run()
+
+            # Verify that year was converted to date range
+            call_kwargs = mock_fetch.call_args[1]
+            assert call_kwargs["start_date"] == "20200101"
+            assert call_kwargs["end_date"] == "20201231"
+
+
+def test_companyprofit_run_with_year_param_int(mock_session, sample_source_data):
+    """Test year parameter as integer"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ", "year": 2020})
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_save_cache"):
+            profit.run()
+
+            call_kwargs = mock_fetch.call_args[1]
+            assert call_kwargs["start_date"] == "20200101"
+            assert call_kwargs["end_date"] == "20201231"
+
+
+def test_companyprofit_run_with_datetime_params(mock_session, sample_source_data):
+    """Test run with datetime objects as parameters"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={
+            "ts_code": "000001.SZ",
+            "start_date": datetime(2020, 1, 1),
+            "end_date": datetime(2020, 12, 31),
+        },
+    )
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_save_cache"):
+            profit.run()
+
+            call_kwargs = mock_fetch.call_args[1]
+            assert call_kwargs["start_date"] == "20200101"
+            assert call_kwargs["end_date"] == "20201231"
+
+
+def test_companyprofit_run_adds_fields_param(mock_session, sample_source_data):
+    """Test that fields parameter is automatically added"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_save_cache"):
+            profit.run()
+
+            call_kwargs = mock_fetch.call_args[1]
+            assert "fields" in call_kwargs
+            assert isinstance(call_kwargs["fields"], str)
+            assert "ts_code" in call_kwargs["fields"]
+
+
+def test_companyprofit_run_sets_cache(mock_session, sample_source_data):
+    """Test that result is saved to cache"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ"},
+        cache=True,
+    )
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data):
+        with patch.object(profit, "_save_cache") as mock_save:
+            profit.run()
+
+            mock_save.assert_called_once()
+
+
+def test_companyprofit_run_calls_transform(mock_session, sample_source_data):
+    """Test that run calls transform method"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data):
+        with patch.object(profit, "transform", return_value=pd.DataFrame()) as mock_transform:
+            with patch.object(profit, "_save_cache"):
+                profit.run()
+
+                mock_transform.assert_called_once()
+
+
+def test_companyprofit_run_uses_income_vip_api(mock_session, sample_source_data):
+    """Test that run method uses the correct API endpoint"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_save_cache"):
+            profit.run()
+
+            # Verify the correct API method is used
+            assert mock_fetch.call_args[1]["api"] == mock_session.connection.income_vip
+
+
+# ============================================================================
+# Integration Tests
+# ============================================================================
+
+
+def test_companyprofit_full_workflow(mock_session, sample_source_data):
+    """Test complete workflow from initialization to result"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ", "start_date": "20200101", "end_date": "20201231"},
+        cache=True,
+    )
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data):
+        with patch.object(profit, "_save_cache"):
+            result = profit.run()
+
+            assert not result.empty
+            assert "code" in result.columns
+            assert "date" in result.columns
+            assert "datecode" in result.columns
+            assert all(result["code"] == "000001.SZ")
+
+
+def test_companyprofit_cache_persistence(mock_session, sample_source_data):
+    """Test that cache persists across runs"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ"},
+        cache=True,
+    )
+
+    # First run - should fetch data
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_load_cache", return_value=None):
+            with patch.object(profit, "_save_cache"):
+                first_result = profit.run()
+                assert mock_fetch.called
+
+    # Second run - should use cache
+    with patch.object(profit, "_fetchall", return_value=sample_source_data) as mock_fetch:
+        with patch.object(profit, "_load_cache", return_value=first_result):
+            profit.run()
+            assert not mock_fetch.called
+
+
+def test_companyprofit_params_identifier_uniqueness(mock_session):
+    """Test that different params create different cache identifiers"""
+    profit1 = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ", "start_date": "20200101"},
+        cache=True,
+    )
+    profit2 = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ", "start_date": "20210101"},
+        cache=True,
+    )
+
+    # Test that different params are properly stored
+    assert profit1.params.get("start_date") == "20200101"
+    assert profit2.params.get("start_date") == "20210101"
+    assert profit1.params.get("start_date") != profit2.params.get("start_date")
+
+
+def test_companyprofit_different_stocks(mock_session, sample_source_data):
+    """Test that different stock codes are handled correctly"""
+    profit1 = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+    profit2 = CompanyProfit(session=mock_session, params={"ts_code": "000002.SZ"})
+
+    assert profit1.params.get("ts_code") != profit2.params.get("ts_code")
+
+
+def test_companyprofit_empty_result_handling(mock_session):
+    """Test handling of empty result from API"""
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+
+    with patch.object(profit, "_fetchall", return_value=pd.DataFrame()):
+        with patch.object(profit, "_save_cache"):
+            result = profit.run()
+
+            assert result.empty
+            assert len(result.columns) == len(TARGET.list_column_names())
+
+
+def test_companyprofit_large_dataset_handling(mock_session):
+    """Test handling of large dataset"""
+    large_data = pd.DataFrame(
+        {
+            "ts_code": ["000001.SZ"] * 1000,
+            "end_date": ["20200331"] * 1000,
+            "ann_date": ["20200425"] * 1000,
+            "f_ann_date": ["20200425"] * 1000,
+            "report_type": ["1"] * 1000,
+            "comp_type": ["1"] * 1000,
+            "end_type": ["12"] * 1000,
+            "update_flag": ["0"] * 1000,
+            "basic_eps": [0.5] * 1000,
+        }
+    )
+
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+
+    with patch.object(profit, "_fetchall", return_value=large_data):
+        with patch.object(profit, "_save_cache"):
+            result = profit.run()
+
+            # After deduplication, should have only 1 row
+            assert len(result) == 1
+
+
+def test_companyprofit_without_cache(mock_session, sample_source_data):
+    """Test that data fetching works without cache"""
+    profit = CompanyProfit(
+        session=mock_session,
+        params={"ts_code": "000001.SZ"},
+        cache=False,
+    )
+
+    with patch.object(profit, "_fetchall", return_value=sample_source_data):
+        result = profit.run()
+
+        assert not result.empty
+        assert profit.cache is None
+
+
+def test_companyprofit_handles_missing_numeric_fields(mock_session):
+    """Test handling of missing numeric values"""
+    data = pd.DataFrame(
+        {
+            "ts_code": ["000001.SZ"],
+            "end_date": ["20200331"],
+            "ann_date": ["20200425"],
+            "f_ann_date": ["20200425"],
+            "report_type": ["1"],
+            "comp_type": ["1"],
+            "end_type": ["12"],
+            "update_flag": ["0"],
+            "basic_eps": [None],
+            "n_income": [None],
+        }
+    )
+
+    profit = CompanyProfit(session=mock_session, params={"ts_code": "000001.SZ"})
+    result = profit.transform(data)
+
+    assert not result.empty
+    assert pd.isna(result["basic_eps"].iloc[0])
+    assert pd.isna(result["n_income"].iloc[0])

xfintech/data/source/tushare/stock/conceptcapflowdc/__init__.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from xfintech.data.source.tushare.stock.conceptcapflowdc.conceptcapflowdc import (
+    ConceptCapflowDC,
+)
+from xfintech.data.source.tushare.stock.conceptcapflowdc.constant import (
+    KEY,
+    NAME,
+    PAGINATE,
+    SOURCE,
+    TARGET,
+)
+
+__all__ = [
+    "ConceptCapflowDC",
+    "KEY",
+    "NAME",
+    "PAGINATE",
+    "SOURCE",
+    "TARGET",
+]