"""Tests for src/atm/levels.py.""" from __future__ import annotations import numpy as np import pytest from atm.config import ( CanaryRegion, ColorSpec, Config, DiscordCfg, ROI, TelegramCfg, YAxisCalib, ) from atm.levels import Levels, LevelsExtractor, LevelsResult from atm.vision import pixel_y_to_price # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- # chart_roi starts at (0,0) so test frames can be exactly (H, W) CHART_ROI = ROI(x=0, y=0, w=600, h=400) CALIB = YAxisCalib(p1_y=0, p1_price=100.0, p2_y=400, p2_price=80.0) # light_red RGB=(255,0,0) → BGR=(0,0,255) # light_green RGB=(0,255,0) → BGR=(0,255,0) RED_BGR: tuple[int, int, int] = (0, 0, 255) GREEN_BGR: tuple[int, int, int] = (0, 255, 0) TOLERANCE = 30.0 def _make_cfg(phaseb_timeout_s: int = 600) -> Config: colors = { "turquoise": ColorSpec(rgb=(0, 255, 255), tolerance=TOLERANCE), "yellow": ColorSpec(rgb=(255, 255, 0), tolerance=TOLERANCE), "dark_green": ColorSpec(rgb=(0, 100, 0), tolerance=TOLERANCE), "dark_red": ColorSpec(rgb=(100, 0, 0), tolerance=TOLERANCE), "light_green": ColorSpec(rgb=(0, 255, 0), tolerance=TOLERANCE), "light_red": ColorSpec(rgb=(255, 0, 0), tolerance=TOLERANCE), "gray": ColorSpec(rgb=(128, 128, 128), tolerance=TOLERANCE), } return Config( window_title="test", dot_roi=ROI(x=0, y=0, w=100, h=50), chart_roi=CHART_ROI, colors=colors, y_axis=CALIB, canary=CanaryRegion( roi=ROI(x=0, y=0, w=50, h=50), baseline_phash="0" * 64, drift_threshold=8, ), discord=DiscordCfg(webhook_url="http://example.com/hook"), telegram=TelegramCfg(bot_token="tok", chat_id="123"), phaseb_timeout_s=phaseb_timeout_s, ) # --------------------------------------------------------------------------- # Frame helpers # --------------------------------------------------------------------------- def _make_chart(*line_specs: tuple[tuple[int, int, int], int]) -> np.ndarray: """Return a 400×600 black BGR frame with horizontal lines. Each spec is (bgr_color, y_position). Lines are painted 3px thick and span the full width so Hough detection is reliable. """ frame = np.zeros((400, 600, 3), dtype=np.uint8) for bgr, y in line_specs: y0 = max(0, y - 1) y1 = min(400, y + 2) frame[y0:y1, :] = bgr return frame # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- def test_three_lines_buy_complete() -> None: """BUY with 3 stable lines → complete after 2 calls, prices correct.""" # red = SL at y=300 (bottom / lowest price for BUY) # green = TP1 at y=200, TP2 at y=100 frame = _make_chart( (RED_BGR, 300), (GREEN_BGR, 200), (GREEN_BGR, 100), ) cfg = _make_cfg() ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0) r1 = ext.step(frame, ts=1.0) assert r1.status == "partial" # first call: not yet stable r2 = ext.step(frame, ts=2.0) assert r2.status == "complete" assert r2.levels is not None assert r2.levels.partial is False expected_tp2 = pixel_y_to_price(100, CALIB) # 95.0 expected_tp1 = pixel_y_to_price(200, CALIB) # 90.0 expected_sl = pixel_y_to_price(300, CALIB) # 85.0 assert r2.levels.tp2 == pytest.approx(expected_tp2, abs=1.0) assert r2.levels.tp1 == pytest.approx(expected_tp1, abs=1.0) assert r2.levels.sl == pytest.approx(expected_sl, abs=1.0) def test_two_lines_partial() -> None: """2 lines → always partial.""" frame = _make_chart((RED_BGR, 300), (GREEN_BGR, 100)) cfg = _make_cfg() ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0) result = ext.step(frame, ts=1.0) assert result.status == "partial" assert result.levels is not None assert result.levels.partial is True def test_zero_lines_waiting() -> None: """No lines in chart → waiting.""" frame = np.zeros((400, 600, 3), dtype=np.uint8) cfg = _make_cfg() ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0) result = ext.step(frame, ts=1.0) assert result.status == "waiting" assert result.levels is None def test_timeout() -> None: """Elapsed > phaseb_timeout_s → timeout regardless of lines.""" frame = np.zeros((400, 600, 3), dtype=np.uint8) cfg = _make_cfg(phaseb_timeout_s=600) ext = LevelsExtractor(cfg, direction="BUY", start_ts=0.0) result = ext.step(frame, ts=700.0) assert result.status == "timeout" assert result.levels is None assert result.elapsed_s == pytest.approx(700.0) def test_sell_direction_assignment() -> None: """SELL: topmost y → SL (highest price), bottom → TP2 (lowest price).""" frame = _make_chart( (RED_BGR, 300), (GREEN_BGR, 200), (GREEN_BGR, 100), ) cfg = _make_cfg() ext = LevelsExtractor(cfg, direction="SELL", start_ts=0.0) ext.step(frame, ts=1.0) # first call, not yet stable r = ext.step(frame, ts=2.0) # second call, stable → complete assert r.status == "complete" assert r.levels is not None expected_sl = pixel_y_to_price(100, CALIB) # 95.0 (topmost = highest price) expected_tp1 = pixel_y_to_price(200, CALIB) # 90.0 expected_tp2 = pixel_y_to_price(300, CALIB) # 85.0 assert r.levels.sl == pytest.approx(expected_sl, abs=1.0) assert r.levels.tp1 == pytest.approx(expected_tp1, abs=1.0) assert r.levels.tp2 == pytest.approx(expected_tp2, abs=1.0)