From e2eb5eb42934d4401f93fef0fc2529e51cf2e1e5 Mon Sep 17 00:00:00 2001
From: "2569718930@qq.com" <2569718930@qq.com>
Date: Thu, 23 Apr 2026 21:15:47 +0800
Subject: [PATCH] feat: implement PolymarketReadOnlyLayer for data collection
 and add comprehensive unit tests

---
 src/data_collection/polymarket_readonly.py | 55 ++++++++++++++++++++--
 tests/test_polymarket_readonly.py          | 42 +++++++++++++++++
 2 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/src/data_collection/polymarket_readonly.py b/src/data_collection/polymarket_readonly.py
index 0977511f..f7767a1f 100644
--- a/src/data_collection/polymarket_readonly.py
+++ b/src/data_collection/polymarket_readonly.py
@@ -1023,7 +1023,12 @@ class PolymarketReadOnlyLayer:
                 preferred_temp=preferred_temp,
             )
 
-        cache_key = f"{city_key}|{target_date}"
+        preferred_temp_key = (
+            f"{float(preferred_temp):.2f}"
+            if preferred_temp is not None
+            else "none"
+        )
+        cache_key = f"{city_key}|{target_date}|{preferred_temp_key}"
         now = time.time()
 
         with self._lock:
@@ -1037,7 +1042,12 @@ class PolymarketReadOnlyLayer:
 
         scored: List[Tuple[float, Dict[str, Any]]] = []
         for market in markets:
-            score = self._score_market(city_key, target_date, market)
+            score = self._score_market(
+                city_key,
+                target_date,
+                market,
+                preferred_temp=preferred_temp,
+            )
             if score <= 0:
                 continue
             scored.append((score, market))
@@ -1046,7 +1056,12 @@ class PolymarketReadOnlyLayer:
         if not scored:
             broader = self._load_markets(active_only=False)
             for market in broader:
-                score = self._score_market(city_key, target_date, market)
+                score = self._score_market(
+                    city_key,
+                    target_date,
+                    market,
+                    preferred_temp=preferred_temp,
+                )
                 if score <= 0:
                     continue
                 scored.append((score, market))
@@ -1196,7 +1211,13 @@ class PolymarketReadOnlyLayer:
 
         return None, f"Specified market_slug not found: {normalized_slug}"
 
-    def _score_market(self, city_key: str, target_date: str, market: Dict[str, Any]) -> float:
+    def _score_market(
+        self,
+        city_key: str,
+        target_date: str,
+        market: Dict[str, Any],
+        preferred_temp: Optional[float] = None,
+    ) -> float:
         city_tokens = CITY_TOKEN_INDEX.get(city_key, [city_key])
         text_parts = [
             market.get("question"),
@@ -1271,6 +1292,32 @@ class PolymarketReadOnlyLayer:
             or 0.0
         )
         score += min(volume / 50000.0, 8.0)
+
+        if preferred_temp is not None:
+            market_temp = self._extract_market_bucket_temp(market)
+            bucket_range = self._extract_market_bucket_range(market)
+            direction = self._extract_market_bucket_direction(market)
+            if market_temp is not None:
+                diff = abs(float(market_temp) - float(preferred_temp))
+                score += max(-40.0, 60.0 - diff * 15.0)
+
+            if bucket_range is not None:
+                lower, upper, _unit = bucket_range
+                contains_preferred = False
+                if upper is not None:
+                    contains_preferred = lower <= float(preferred_temp) <= upper
+                elif direction == "above":
+                    contains_preferred = float(preferred_temp) >= lower
+                elif direction == "below":
+                    contains_preferred = float(preferred_temp) <= lower
+                else:
+                    contains_preferred = abs(float(preferred_temp) - lower) <= 0.51
+
+                if contains_preferred:
+                    if upper is not None or direction == "exact":
+                        score += 28.0
+                    else:
+                        score += 18.0
         return score
 
     def _is_temperature_market(self, market: Dict[str, Any]) -> bool:
diff --git a/tests/test_polymarket_readonly.py b/tests/test_polymarket_readonly.py
index c34d41d0..42540359 100644
--- a/tests/test_polymarket_readonly.py
+++ b/tests/test_polymarket_readonly.py
@@ -442,3 +442,45 @@ def test_build_top_temperature_buckets_dedupes_same_temperature():
         "highest-temperature-in-ankara-on-march-12-2026-14c-or-higher|yes"
     )
     assert all(not str(row.get("label") or "").startswith("<=") for row in rows)
+
+
+def test_find_primary_market_prefers_preferred_temperature_and_cache_key():
+    layer = PolymarketReadOnlyLayer()
+    markets = [
+        {
+            "slug": "highest-temperature-in-madrid-on-april-23-2026-22corbelow",
+            "question": "Will the highest temperature in Madrid be 22C or below on April 23?",
+            "volumeNum": 900000,
+            "active": True,
+            "closed": False,
+            "acceptingOrders": True,
+            "enableOrderBook": True,
+        },
+        {
+            "slug": "highest-temperature-in-madrid-on-april-23-2026-27c",
+            "question": "Will the highest temperature in Madrid be 27C on April 23?",
+            "volumeNum": 1000,
+            "active": True,
+            "closed": False,
+            "acceptingOrders": True,
+            "enableOrderBook": True,
+        },
+    ]
+
+    layer._load_markets = lambda active_only=True: markets
+
+    selected_27, reason_27 = layer._find_primary_market(
+        "madrid",
+        "2026-04-23",
+        preferred_temp=27.0,
+    )
+    selected_22, reason_22 = layer._find_primary_market(
+        "madrid",
+        "2026-04-23",
+        preferred_temp=22.0,
+    )
+
+    assert reason_27 is None
+    assert reason_22 is None
+    assert selected_27["slug"] == "highest-temperature-in-madrid-on-april-23-2026-27c"
+    assert selected_22["slug"] == "highest-temperature-in-madrid-on-april-23-2026-22corbelow"