mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 15:34:16 +08:00
287 lines
11 KiB
Python
287 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import UTC, date, datetime, timedelta
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from app.services import heavy_job_limiter, pipeline_jobs, preferences
|
|
from app.services.heavy_job_limiter import HeavyJobLimiter
|
|
from app.services.pipeline_jobs import JobCancelledError, JobStore, run_with_capacity
|
|
from app.tickflow.repository import KlineRepository
|
|
|
|
|
|
def wait_for(predicate, timeout=2):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return
|
|
time.sleep(0.005)
|
|
raise AssertionError("condition did not become true")
|
|
|
|
|
|
def waiting(store, jid):
|
|
log = store.get(jid)["log"]
|
|
return bool(log) and log[-1]["stage"] == "init"
|
|
|
|
|
|
@pytest.fixture
|
|
def capacity(monkeypatch, tmp_path):
|
|
monkeypatch.setattr(preferences, "load", lambda: {})
|
|
monkeypatch.setattr(pipeline_jobs, "_CANCEL_FLAGS", {})
|
|
monkeypatch.setattr(pipeline_jobs, "_run_slot_owner", None)
|
|
store = JobStore(store_dir=tmp_path / "jobs")
|
|
limiter = HeavyJobLimiter(capacity=2, cancel_poll_interval=0.005)
|
|
monkeypatch.setattr(pipeline_jobs, "job_store", store)
|
|
monkeypatch.setattr(heavy_job_limiter, "shared_heavy_job_limiter", limiter)
|
|
return store, limiter
|
|
|
|
|
|
def test_queued_pipeline_is_cancellable_and_not_reaped(capacity):
|
|
store, limiter = capacity
|
|
jid, _ = store.create(timeout_s=1)
|
|
called = threading.Event()
|
|
assert limiter.acquire("normal", timeout=0)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(run_with_capacity, jid, called.set)
|
|
try:
|
|
wait_for(lambda: waiting(store, jid))
|
|
assert store.get(jid)["status"] == "pending"
|
|
stale = (datetime.now(UTC) - timedelta(hours=2)).isoformat()
|
|
store._active_jobs[jid]["started_at"] = stale
|
|
store._active_jobs[jid]["last_progress_at"] = stale
|
|
store.reap_stale()
|
|
assert store.get(jid)["status"] == "pending"
|
|
store.terminate(jid, "cancelled in queue")
|
|
with pytest.raises(JobCancelledError):
|
|
future.result(timeout=1)
|
|
assert not called.is_set()
|
|
assert limiter.in_use == 1
|
|
finally:
|
|
limiter.release("normal")
|
|
assert limiter.in_use == 0
|
|
|
|
|
|
def test_cancel_does_not_release_running_worker_capacity(capacity):
|
|
store, limiter = capacity
|
|
jid, _ = store.create()
|
|
entered, finish = threading.Event(), threading.Event()
|
|
|
|
def work():
|
|
entered.set()
|
|
assert finish.wait(2)
|
|
raise JobCancelledError(jid)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(run_with_capacity, jid, work)
|
|
try:
|
|
assert entered.wait(1)
|
|
store.terminate(jid, "cancelled while computing")
|
|
assert store.get(jid)["status"] == "failed"
|
|
assert limiter.in_use == 2
|
|
assert not limiter.acquire("normal", timeout=0)
|
|
finally:
|
|
finish.set()
|
|
with pytest.raises(JobCancelledError):
|
|
future.result(timeout=1)
|
|
assert limiter.in_use == 0
|
|
|
|
|
|
def test_cache_refresh_queues_and_can_nest_in_pipeline(capacity, monkeypatch):
|
|
store, limiter = capacity
|
|
repo = object.__new__(KlineRepository)
|
|
calls = []
|
|
monkeypatch.setattr(repo, "_refresh_enriched_impl", lambda: calls.append(limiter.in_use))
|
|
assert limiter.acquire("normal", timeout=0)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(repo._refresh_enriched)
|
|
try:
|
|
wait_for(lambda: bool(limiter._waiters))
|
|
assert calls == []
|
|
finally:
|
|
limiter.release("normal")
|
|
future.result(timeout=1)
|
|
jid, _ = store.create()
|
|
run_with_capacity(jid, repo._refresh_enriched)
|
|
assert calls == [2, 2]
|
|
assert limiter.in_use == 0
|
|
|
|
|
|
@pytest.mark.parametrize("kind", ["signal", "factor", "factor_batch"])
|
|
def test_in_process_backtests_share_capacity(capacity, monkeypatch, kind):
|
|
from app.backtest.factor import FactorBacktestService
|
|
from app.services.backtest import BacktestService
|
|
|
|
_, limiter = capacity
|
|
cls = BacktestService if kind == "signal" else FactorBacktestService
|
|
service = object.__new__(cls)
|
|
method = "run_batch" if kind == "factor_batch" else "run"
|
|
received = []
|
|
config = object()
|
|
|
|
def compute(actual, **kwargs):
|
|
assert actual is config
|
|
assert limiter.in_use == 2
|
|
received.append(True)
|
|
return "unchanged result"
|
|
|
|
monkeypatch.setattr(service, "_" + method, compute)
|
|
assert limiter.acquire("normal", timeout=0)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(getattr(service, method), config)
|
|
try:
|
|
wait_for(lambda: bool(limiter._waiters))
|
|
assert received == []
|
|
finally:
|
|
limiter.release("normal")
|
|
assert future.result(timeout=1) == "unchanged result"
|
|
# Auto-mining evaluates factors within its existing two-slot reservation.
|
|
with limiter.slot("mining"):
|
|
assert getattr(service, method)(config) == "unchanged result"
|
|
assert limiter.in_use == 0
|
|
|
|
|
|
@pytest.mark.parametrize("endpoint,body,module_name,function_name,result", [
|
|
("/api/pipeline/run", {}, "app.jobs.daily_pipeline", "run_now", {}),
|
|
("/api/kline/extend_history", {"value": 1, "unit": "month"},
|
|
"app.services.extend_history", "run_extend_history", {}),
|
|
("/api/kline/repair_daily", {"start_date": "2026-01-01"},
|
|
"app.services.repair_daily", "run_repair_daily", {}),
|
|
("/api/kline/rebuild_enriched", {}, "app.indicators.pipeline", "run_pipeline", 3),
|
|
("/api/kline/sync_minute", {"days": 1},
|
|
"app.services.kline_sync", "sync_and_persist_minute", 3),
|
|
])
|
|
def test_api_jobs_wait_before_computing(
|
|
capacity, monkeypatch, tmp_path, endpoint, body, module_name, function_name, result,
|
|
):
|
|
import importlib
|
|
|
|
from app.api import kline, pipeline
|
|
from app.jobs import daily_pipeline
|
|
|
|
store, limiter = capacity
|
|
entered, finish = threading.Event(), threading.Event()
|
|
|
|
def compute(*args, **kwargs):
|
|
assert limiter.in_use == 2
|
|
entered.set()
|
|
assert finish.wait(3)
|
|
return result
|
|
|
|
monkeypatch.setattr(importlib.import_module(module_name), function_name, compute)
|
|
monkeypatch.setattr(pipeline, "job_store", store)
|
|
monkeypatch.setattr("app.tickflow.pools.get_pool", lambda *args: [])
|
|
monkeypatch.setattr(daily_pipeline, "_refresh_single_view", lambda *args: None)
|
|
app = FastAPI()
|
|
app.include_router(pipeline.router)
|
|
app.include_router(kline.router)
|
|
app.state.repo = SimpleNamespace(
|
|
store=SimpleNamespace(data_dir=tmp_path),
|
|
db=SimpleNamespace(execute=lambda *args: None),
|
|
refresh_cache=lambda: None,
|
|
get_index_symbol_set=lambda: set(),
|
|
)
|
|
app.state.capabilities = SimpleNamespace(has=lambda _: True)
|
|
assert limiter.acquire("normal", timeout=0)
|
|
released = False
|
|
with TestClient(app) as client:
|
|
try:
|
|
response = client.post(endpoint, json=body)
|
|
assert response.status_code == 200, response.text
|
|
jid = response.json()["job_id"]
|
|
wait_for(lambda: waiting(store, jid))
|
|
assert not entered.is_set()
|
|
assert client.get(f"/api/pipeline/jobs/{jid}").json()["status"] == "pending"
|
|
limiter.release("normal")
|
|
released = True
|
|
assert entered.wait(2)
|
|
assert not limiter.acquire("normal", timeout=0)
|
|
finish.set()
|
|
wait_for(lambda: store.get(jid)["status"] in ("succeeded", "failed"))
|
|
assert store.get(jid)["status"] == "succeeded", store.get(jid)
|
|
finally:
|
|
finish.set()
|
|
if not released:
|
|
limiter.release("normal")
|
|
wait_for(lambda: limiter.in_use == 0)
|
|
|
|
|
|
@pytest.mark.parametrize("fail", [False, True])
|
|
def test_manual_pipeline_holds_capacity_through_refresh_off_event_loop(capacity, monkeypatch, fail):
|
|
from app.api import pipeline
|
|
|
|
store, limiter = capacity
|
|
refreshing, finish = threading.Event(), threading.Event()
|
|
|
|
def compute(*args, **kwargs):
|
|
if fail:
|
|
raise ValueError("partial pipeline failure")
|
|
return {"rows": 3}
|
|
|
|
def refresh():
|
|
assert limiter.in_use == 2
|
|
refreshing.set()
|
|
assert finish.wait(3)
|
|
|
|
monkeypatch.setattr(pipeline, "job_store", store)
|
|
monkeypatch.setattr(pipeline.daily_pipeline, "run_now", compute)
|
|
app = FastAPI()
|
|
app.include_router(pipeline.router)
|
|
app.state.repo = SimpleNamespace(refresh_cache=refresh)
|
|
app.state.capabilities = object()
|
|
with TestClient(app) as client:
|
|
try:
|
|
jid = client.post("/api/pipeline/run").json()["job_id"]
|
|
assert refreshing.wait(1)
|
|
assert not limiter.acquire("normal", timeout=0)
|
|
# This request must complete while the worker is blocked in refresh.
|
|
assert client.get(f"/api/pipeline/jobs/{jid}").json()["status"] == "running"
|
|
finish.set()
|
|
wait_for(lambda: store.get(jid)["status"] in ("succeeded", "failed"))
|
|
assert store.get(jid)["status"] == ("failed" if fail else "succeeded")
|
|
finally:
|
|
finish.set()
|
|
assert limiter.in_use == 0
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ["scheduled", "integrity"])
|
|
def test_background_pipeline_entries_share_capacity(capacity, monkeypatch, entry):
|
|
from app.jobs import daily_pipeline
|
|
from app.services import data_integrity, repair_daily
|
|
|
|
store, limiter = capacity
|
|
called = threading.Event()
|
|
|
|
def compute(*args, **kwargs):
|
|
assert limiter.in_use == 2
|
|
called.set()
|
|
return {}
|
|
|
|
monkeypatch.setattr(repair_daily, "run_repair_daily", compute)
|
|
assert limiter.acquire("normal", timeout=0)
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
if entry == "scheduled":
|
|
future = pool.submit(daily_pipeline._run_tracked, compute, "test")
|
|
else:
|
|
state = SimpleNamespace(repo=object(), capabilities=SimpleNamespace(has=lambda _: True))
|
|
data_integrity.launch_integrity_repair(state, date(2026, 1, 1), "test")
|
|
future = None
|
|
try:
|
|
wait_for(lambda: store.active_id() is not None)
|
|
jid = store.active_id()
|
|
wait_for(lambda: waiting(store, jid))
|
|
assert store.get(jid)["status"] == "pending"
|
|
assert not called.is_set()
|
|
finally:
|
|
limiter.release("normal")
|
|
wait_for(lambda: store.get(jid)["status"] == "succeeded")
|
|
if future is not None:
|
|
assert future.result(timeout=1)
|
|
assert called.is_set()
|
|
assert limiter.in_use == 0
|