feat(portfolio): LocalUnifiedProvider 成份股并集治偏差(Task2)

2026-07-23 08:07:37 +08:00
parent 4b4dca2f44
commit b46a0c26c8
2 changed files with 102 additions and 0 deletions
@@ -280,3 +280,63 @@ class TestGetPrice:
        )
        assert len(df) == 2
        assert set(df["code"]) == {"600519.XSHG", "000001.XSHE"}
+
+
+# ======================== Task 2: get_index_stocks ========================
+def _make_constituent_db(tmp_path):
+    db = tmp_path / "t.db"
+    c = sqlite3.connect(str(db))
+    c.execute(
+        "CREATE TABLE constituent_unified(index_code TEXT, code TEXT, code_name TEXT, "
+        "source TEXT, in_current INT, was_removed INT)"
+    )
+    c.executemany(
+        "INSERT INTO constituent_unified VALUES(?,?,?,?,?,?)",
+        [
+            ("000300", "600519", "贵州茅台", "baostock", 1, 0),
+            ("000300", "000001", "平安银行", "baostock", 1, 0),
+            ("000300", "600811", "退市股", "baostock", 0, 1),  # 被踢
+        ],
+    )
+    c.commit()
+    c.close()
+    return db
+
+
+class TestGetIndexStocks:
+    def test_union_of_current_and_removed(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        stocks = p.get_index_stocks("000300.XSHG", "2020-01-01")
+        # 并集含被踢(was_removed=1)
+        assert "600519.XSHG" in stocks
+        assert "000001.XSHE" in stocks
+        assert "600811.XSHG" in stocks  # 6 开头 → SSE
+        assert len(stocks) == 3
+
+    def test_date_param_ignored_union_model(self, tmp_path):
+        # 并集模型 — date 参数不报错不过滤
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        s1 = p.get_index_stocks("000300", "2010-01-01")
+        s2 = p.get_index_stocks("000300", "2024-12-31")
+        assert set(s1) == set(s2)
+
+    def test_get_constituent_is_alias(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        a = p.get_index_stocks("000300.XSHG", "2020-01-01")
+        b = p.get_constituent("000300", "2020-01-01")
+        assert a == b
+
+    def test_index_not_found_returns_empty(self, tmp_path):
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        assert p.get_index_stocks("999999.XSHG", "2024-01-01") == []
+
+    def test_pure_digit_index_code(self, tmp_path):
+        # 纯数字 index_symbol 也能查
+        db = _make_constituent_db(tmp_path)
+        p = LocalUnifiedProvider({"db_path": str(db), "data_dir": str(tmp_path)})
+        stocks = p.get_index_stocks("000300", "2020-01-01")
+        assert len(stocks) == 3