Blame view

tests/test_process_products_batching.py 2.8 KB
be3f0d46   tangwang   /indexer/enrich-c...
1
2
3
4
  from __future__ import annotations
  
  from typing import Any, Dict, List
  
6f7840cf   tangwang   refactor: rename ...
5
  import indexer.product_enrich as process_products
be3f0d46   tangwang   /indexer/enrich-c...
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
  
  
  def _mk_products(n: int) -> List[Dict[str, str]]:
      return [{"id": str(i), "title": f"title-{i}"} for i in range(n)]
  
  
  def test_analyze_products_caps_batch_size_to_20(monkeypatch):
      monkeypatch.setattr(process_products, "API_KEY", "fake-key")
      seen_batch_sizes: List[int] = []
  
      def _fake_process_batch(batch_data: List[Dict[str, str]], batch_num: int, target_lang: str = "zh"):
          seen_batch_sizes.append(len(batch_data))
          return [
              {
                  "id": item["id"],
                  "lang": target_lang,
                  "title_input": item["title"],
                  "title": "",
                  "category_path": "",
                  "tags": "",
                  "target_audience": "",
                  "usage_scene": "",
                  "season": "",
                  "key_attributes": "",
                  "material": "",
                  "features": "",
                  "selling_points": "",
                  "anchor_text": "",
              }
              for item in batch_data
          ]
  
      monkeypatch.setattr(process_products, "process_batch", _fake_process_batch)
      monkeypatch.setattr(process_products, "_set_cached_anchor_result", lambda *args, **kwargs: None)
  
      out = process_products.analyze_products(
          products=_mk_products(45),
          target_lang="zh",
          batch_size=200,
          tenant_id="162",
      )
  
      assert len(out) == 45
      assert seen_batch_sizes == [20, 20, 5]
  
  
  def test_analyze_products_uses_min_batch_size_1(monkeypatch):
      monkeypatch.setattr(process_products, "API_KEY", "fake-key")
      seen_batch_sizes: List[int] = []
  
      def _fake_process_batch(batch_data: List[Dict[str, str]], batch_num: int, target_lang: str = "zh"):
          seen_batch_sizes.append(len(batch_data))
          return [
              {
                  "id": item["id"],
                  "lang": target_lang,
                  "title_input": item["title"],
                  "title": "",
                  "category_path": "",
                  "tags": "",
                  "target_audience": "",
                  "usage_scene": "",
                  "season": "",
                  "key_attributes": "",
                  "material": "",
                  "features": "",
                  "selling_points": "",
                  "anchor_text": "",
              }
              for item in batch_data
          ]
  
      monkeypatch.setattr(process_products, "process_batch", _fake_process_batch)
      monkeypatch.setattr(process_products, "_set_cached_anchor_result", lambda *args, **kwargs: None)
  
      out = process_products.analyze_products(
          products=_mk_products(3),
          target_lang="zh",
          batch_size=0,
          tenant_id="162",
      )
  
      assert len(out) == 3
      assert seen_batch_sizes == [1, 1, 1]