Image vector embedding
This commit is contained in:
sriram
@@ -85,6 +85,12 @@ os.environ["ACTIVE_BRANDS"] = ""
|
||||
# turns it on explicitly, the same arrangement ACTIVE_BRANDS has above.
|
||||
os.environ["AUTO_ENRICH_ON_UPLOAD"] = "false"
|
||||
|
||||
# Same arrangement for the image-vector worker: every call to
|
||||
# `upsert_brand_products` would otherwise start a thread that opens a
|
||||
# database connection and downloads product images. Off here; the hook and
|
||||
# the worker are covered explicitly in tests/test_image_vector.py.
|
||||
os.environ["ENABLE_IMAGE_VECTORS"] = "false"
|
||||
|
||||
# Auth is set unconditionally (not setdefault): the suite asserts on the real
|
||||
# guards, so it must never inherit a developer's AUTH_ENABLED=false.
|
||||
os.environ["AUTH_ENABLED"] = "true"
|
||||
|
||||
649
tests/test_image_vector.py
Normal file
649
tests/test_image_vector.py
Normal file
@@ -0,0 +1,649 @@
|
||||
"""img_vector: the pixel thumbnail of each product's primary image.
|
||||
|
||||
Four things have to hold, and each has broken independently for a sibling
|
||||
column before, so each gets its own tests here:
|
||||
|
||||
1. The vector is what the card shows: `image_url`, else `image_urls[0]`,
|
||||
EXIF-rotated, flattened onto white, 32x32 RGB - and anything Pillow cannot
|
||||
decode is None, never an exception (pytest runs warnings as errors).
|
||||
2. The columns exist on every brand table via `_ensure_columns`, and a server
|
||||
that refuses `vector(3072)` loses only that column, not the write.
|
||||
3. The upsert never names them (the nutrition_score rule), and the product
|
||||
readers never select them - the browse endpoint pulls the whole catalog.
|
||||
4. The hook is off the write path: it enqueues and returns, it is a no-op
|
||||
when disabled, and a full queue is a log line, not an error.
|
||||
|
||||
No database is involved anywhere; cursors are recorders.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import queue
|
||||
import re
|
||||
import threading
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from app.services import image_vector as iv
|
||||
from app.services import vector_store
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _png(im: Image.Image) -> bytes:
|
||||
buf = io.BytesIO()
|
||||
im.save(buf, "PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _jpeg(im: Image.Image, **kw) -> bytes:
|
||||
buf = io.BytesIO()
|
||||
im.save(buf, "JPEG", quality=95, **kw)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _pixel(vec: List[int], x: int, y: int) -> tuple:
|
||||
i = (y * iv.IMG_VECTOR_SIZE + x) * 3
|
||||
return tuple(vec[i:i + 3])
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_column_cache():
|
||||
vector_store._invalidate_product_columns_cache()
|
||||
yield
|
||||
vector_store._invalidate_product_columns_cache()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. bytes -> vector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_a_solid_png_becomes_3072_values_of_that_colour():
|
||||
vec = iv.pixel_vector(_png(Image.new("RGB", (200, 120), (255, 0, 0))))
|
||||
|
||||
assert vec is not None
|
||||
assert len(vec) == iv.IMG_VECTOR_DIM == 3072
|
||||
assert set(vec[0::3]) == {255}
|
||||
assert set(vec[1::3]) == {0}
|
||||
assert set(vec[2::3]) == {0}
|
||||
|
||||
|
||||
def test_a_solid_jpeg_is_within_lossy_tolerance():
|
||||
vec = iv.pixel_vector(_jpeg(Image.new("RGB", (300, 300), (40, 120, 200))))
|
||||
|
||||
assert vec is not None and len(vec) == 3072
|
||||
for channel, expected in enumerate((40, 120, 200)):
|
||||
assert all(abs(v - expected) <= 3 for v in vec[channel::3])
|
||||
|
||||
|
||||
def test_exif_orientation_is_applied_because_the_browser_applies_it():
|
||||
"""Left half red, right half blue, tagged 'rotate 90 CW to display'.
|
||||
|
||||
Without the transpose, the bottom-left pixel is red (it is the left half).
|
||||
With it, the left half has become the top half and bottom-left is blue.
|
||||
"""
|
||||
im = Image.new("RGB", (64, 32), (255, 0, 0))
|
||||
im.paste((0, 0, 255), (32, 0, 64, 32))
|
||||
exif = Image.Exif()
|
||||
exif[0x0112] = 6
|
||||
vec = iv.pixel_vector(_jpeg(im, exif=exif.tobytes()))
|
||||
|
||||
assert vec is not None
|
||||
r, g, b = _pixel(vec, 0, iv.IMG_VECTOR_SIZE - 1)
|
||||
assert b > 200 and r < 60, (r, g, b)
|
||||
r, g, b = _pixel(vec, 0, 0)
|
||||
assert r > 200 and b < 60, (r, g, b)
|
||||
|
||||
|
||||
def test_transparency_is_flattened_onto_white_not_black():
|
||||
im = Image.new("RGBA", (50, 50), (0, 0, 0, 0)) # fully transparent
|
||||
vec = iv.pixel_vector(_png(im))
|
||||
|
||||
assert vec is not None
|
||||
assert set(vec) == {255}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["P", "L", "LA", "CMYK", "I;16"])
|
||||
def test_every_pillow_mode_a_product_photo_could_arrive_in_decodes(mode):
|
||||
base = Image.new("RGB", (40, 40), (10, 200, 30))
|
||||
im = base.convert(mode) if mode != "P" else base.quantize(16)
|
||||
buf = io.BytesIO()
|
||||
im.save(buf, "TIFF" if mode in ("CMYK", "I;16") else "PNG")
|
||||
|
||||
vec = iv.pixel_vector(buf.getvalue())
|
||||
|
||||
assert vec is not None and len(vec) == 3072
|
||||
assert all(0 <= v <= 255 for v in vec)
|
||||
|
||||
|
||||
def test_the_first_frame_of_an_animated_gif_is_used():
|
||||
frames = [Image.new("P", (20, 20), c) for c in (1, 2, 3)]
|
||||
buf = io.BytesIO()
|
||||
frames[0].save(buf, "GIF", save_all=True, append_images=frames[1:])
|
||||
|
||||
assert iv.pixel_vector(buf.getvalue()) is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data", [b"", b"not an image", b"\x89PNG\r\n\x1a\n" + b"\x00" * 40])
|
||||
def test_undecodable_bytes_are_none_not_an_exception(data):
|
||||
assert iv.pixel_vector(data) is None
|
||||
|
||||
|
||||
def test_a_truncated_file_is_none():
|
||||
noisy = Image.effect_noise((64, 64), 50).convert("RGB")
|
||||
whole = _png(noisy)
|
||||
assert len(whole) > 400, "need a file big enough to cut"
|
||||
|
||||
assert iv.pixel_vector(whole[:200]) is None
|
||||
|
||||
|
||||
def test_a_header_claiming_too_many_pixels_is_refused_before_decoding(monkeypatch):
|
||||
monkeypatch.setattr(iv, "IMAGE_VECTOR_MAX_PIXELS", 100)
|
||||
|
||||
assert iv.pixel_vector(_png(Image.new("RGB", (20, 20)))) is None
|
||||
assert iv.pixel_vector(_png(Image.new("RGB", (10, 10)))) is not None
|
||||
|
||||
|
||||
def test_to_pg_is_the_same_text_form_the_upsert_uses_for_embedding():
|
||||
assert iv.to_pg([0, 128, 255]) == "[0,128,255]"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1b. which image
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_primary_is_image_url_then_image_urls_zero_like_the_card():
|
||||
assert iv.primary_image_url({"image_url": "https://a/x.jpg", "image_urls": ["https://b/y.jpg"]}) == "https://a/x.jpg"
|
||||
assert iv.primary_image_url({"image_url": "", "image_urls": ["https://b/y.jpg", "https://c/z.jpg"]}) == "https://b/y.jpg"
|
||||
assert iv.primary_image_url({"image_url": None, "image_urls": None}) is None
|
||||
assert iv.primary_image_url({}) is None
|
||||
|
||||
|
||||
def test_primary_is_normalised_the_way_the_card_fallbacks_are():
|
||||
assert iv.primary_image_url({"image_url": "//cdn.example/x.jpg"}) == "https://cdn.example/x.jpg"
|
||||
dead = next(iter(vector_store._DEAD_IMAGE_HOSTS))
|
||||
assert iv.primary_image_url({"image_url": f"https://{dead}/daily/x.jpg"}) is None
|
||||
assert iv.primary_image_url({"image_url": "not a url"}) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1c. URL -> bytes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _Resp:
|
||||
def __init__(self, status: int, body: bytes, content_type: str = "image/jpeg"):
|
||||
self.status_code = status
|
||||
self.headers = {"content-type": content_type}
|
||||
self._body = body
|
||||
self.closed = False
|
||||
|
||||
def iter_content(self, chunk_size=1):
|
||||
for i in range(0, len(self._body), chunk_size):
|
||||
yield self._body[i:i + chunk_size]
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
def _patch_requests(monkeypatch, responses: List[_Resp]):
|
||||
calls: List[Dict[str, Any]] = []
|
||||
|
||||
def fake_get(url, headers=None, timeout=None, stream=False):
|
||||
calls.append({"url": url, "headers": headers})
|
||||
return responses.pop(0)
|
||||
|
||||
import requests
|
||||
monkeypatch.setattr(requests, "get", fake_get)
|
||||
return calls
|
||||
|
||||
|
||||
def test_a_403_with_referer_is_retried_without_one(monkeypatch):
|
||||
body = b"x" * (iv.MIN_IMAGE_BYTES + 10)
|
||||
calls = _patch_requests(monkeypatch, [_Resp(403, b""), _Resp(200, body)])
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1.jpg") == body
|
||||
assert "Referer" in calls[0]["headers"]
|
||||
assert "Referer" not in calls[1]["headers"]
|
||||
|
||||
|
||||
def test_a_download_past_the_byte_cap_is_abandoned(monkeypatch):
|
||||
monkeypatch.setattr(iv, "IMAGE_VECTOR_MAX_BYTES", 1000)
|
||||
_patch_requests(monkeypatch, [_Resp(200, b"x" * 5000)])
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1.jpg") is None
|
||||
|
||||
|
||||
def test_a_placeholder_below_min_image_bytes_is_rejected(monkeypatch):
|
||||
_patch_requests(monkeypatch, [_Resp(200, b"x" * 10), _Resp(200, b"x" * 10)])
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1.jpg") is None
|
||||
|
||||
|
||||
def test_a_non_image_content_type_without_an_image_extension_is_rejected(monkeypatch):
|
||||
_patch_requests(monkeypatch, [_Resp(200, b"<html>" * 1000, "text/html")])
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1") is None
|
||||
|
||||
|
||||
def test_text_html_is_rejected_even_when_the_url_ends_in_png(monkeypatch):
|
||||
"""uat.amul.com redirects every dead path to its homepage. The .png in the
|
||||
URL is no evidence; the server's own content-type is."""
|
||||
_patch_requests(monkeypatch, [_Resp(200, b"<html>" * 1000, "text/html; charset=utf-8")])
|
||||
|
||||
assert iv.download_image_bytes("https://uat.amul.com/files/products/amul-ghee1.png") is None
|
||||
|
||||
|
||||
def test_an_unlabelled_body_with_an_image_extension_is_accepted(monkeypatch):
|
||||
body = b"x" * (iv.MIN_IMAGE_BYTES + 10)
|
||||
_patch_requests(monkeypatch, [_Resp(200, body, "application/octet-stream")])
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1.jpg") == body
|
||||
|
||||
|
||||
def test_a_raised_request_is_none_not_an_exception(monkeypatch):
|
||||
import requests
|
||||
|
||||
def boom(*a, **k):
|
||||
raise requests.ConnectionError("no route")
|
||||
|
||||
monkeypatch.setattr(requests, "get", boom)
|
||||
|
||||
assert iv.download_image_bytes("https://cdn.example/p/1.jpg") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. the columns exist, and an old pgvector loses only the vector
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MigrationCursor:
|
||||
def __init__(self, refuse: str = ""):
|
||||
self.statements: List[str] = []
|
||||
self._refuse = refuse
|
||||
|
||||
def execute(self, sql, params=None):
|
||||
text = " ".join(str(sql).split())
|
||||
self.statements.append(text)
|
||||
if self._refuse and self._refuse in text and "ADD COLUMN" in text:
|
||||
raise RuntimeError('type "vector(3072)" does not exist')
|
||||
|
||||
def fetchall(self):
|
||||
return []
|
||||
|
||||
|
||||
def test_the_migration_adds_both_columns_to_an_existing_table():
|
||||
cur = MigrationCursor()
|
||||
|
||||
vector_store._ensure_columns(cur, "brand_cadbury")
|
||||
|
||||
assert "ALTER TABLE brand_cadbury ADD COLUMN IF NOT EXISTS img_vector vector(3072)" in cur.statements
|
||||
assert "ALTER TABLE brand_cadbury ADD COLUMN IF NOT EXISTS img_vector_src TEXT" in cur.statements
|
||||
|
||||
|
||||
def test_a_server_that_refuses_the_vector_type_still_gets_every_other_column():
|
||||
cur = MigrationCursor(refuse="img_vector vector(3072)")
|
||||
|
||||
vector_store._ensure_columns(cur, "brand_cadbury") # must not raise
|
||||
|
||||
assert "ALTER TABLE brand_cadbury ADD COLUMN IF NOT EXISTS img_vector_src TEXT" in cur.statements
|
||||
assert "ALTER TABLE brand_cadbury ADD COLUMN IF NOT EXISTS embedding vector(384)" in cur.statements
|
||||
|
||||
|
||||
def test_the_create_table_ddl_carries_them_and_does_not_index_the_vector():
|
||||
ddl = vector_store.get_brand_table_ddl("Cadbury")
|
||||
|
||||
assert "img_vector vector(3072)" in ddl
|
||||
assert "img_vector_src TEXT" in ddl
|
||||
for line in ddl.splitlines():
|
||||
if "CREATE INDEX" in line:
|
||||
assert "img_vector" not in line, line
|
||||
|
||||
|
||||
def test_the_migration_script_can_still_read_col_defs_as_literals():
|
||||
from scripts.migrate_brand_schema import _declared_types
|
||||
|
||||
declared = _declared_types()
|
||||
|
||||
assert declared["img_vector"] == "vector(3072)"
|
||||
assert declared["img_vector_src"] == "TEXT"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. the upsert never names them; the readers never select them
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _insert_statement() -> str:
|
||||
import inspect
|
||||
body = inspect.getsource(vector_store.upsert_brand_products)
|
||||
match = re.search(r"INSERT INTO \{table_name\}.*?updated_at = CURRENT_TIMESTAMP", body, re.S)
|
||||
assert match
|
||||
return " ".join(match.group(0).split())
|
||||
|
||||
|
||||
def test_the_insert_does_not_name_the_vector_columns():
|
||||
"""A column the statement never names is a column it cannot damage. Every
|
||||
writer reaching the upsert builds its dict from a sheet, the seed JSON or
|
||||
a read-back row; none can carry a pixel vector, so naming the column here
|
||||
would only ever NULL it on the next re-seed."""
|
||||
statement = _insert_statement()
|
||||
|
||||
assert "img_vector" not in statement
|
||||
|
||||
|
||||
def test_the_write_hook_runs_after_the_write_and_cannot_raise_into_it():
|
||||
import inspect
|
||||
body = inspect.getsource(vector_store.upsert_brand_products)
|
||||
|
||||
hook = body.index("_schedule_image_vectors(brand, expected_ids)")
|
||||
assert body.index("return persisted") > hook > body.index("invalidate_brand_overview_cache()")
|
||||
assert "except Exception" in body[body.rindex("try:", 0, hook):hook + 200]
|
||||
|
||||
|
||||
class ColumnsCursor:
|
||||
"""Answers the information_schema probe, the table-exists probe and one
|
||||
product SELECT; records every statement."""
|
||||
|
||||
def __init__(self, tables: Dict[str, List[str]]):
|
||||
self.tables = tables
|
||||
self.statements: List[str] = []
|
||||
self.description = None
|
||||
self._pending: Any = None
|
||||
|
||||
def execute(self, sql, params=None):
|
||||
text = " ".join(str(sql).split())
|
||||
self.statements.append(text)
|
||||
if "information_schema.columns" in text:
|
||||
self._pending = [(t, c) for t, cols in self.tables.items() for c in cols]
|
||||
elif "information_schema.tables" in text:
|
||||
self._pending = (True,)
|
||||
else:
|
||||
self.description = [("id",), ("product_name",)]
|
||||
self._pending = []
|
||||
|
||||
def fetchall(self):
|
||||
return self._pending if isinstance(self._pending, list) else []
|
||||
|
||||
def fetchone(self):
|
||||
return self._pending if isinstance(self._pending, tuple) else None
|
||||
|
||||
|
||||
def test_product_columns_drops_both_vectors_and_keeps_the_provenance_url():
|
||||
cur = ColumnsCursor({"brand_x": ["id", "embedding", "img_vector", "img_vector_src", "size"]})
|
||||
|
||||
assert vector_store._product_columns(cur, "brand_x") == '"id", "img_vector_src", "size"'
|
||||
|
||||
|
||||
def test_an_unknown_table_falls_back_to_star_rather_than_failing():
|
||||
cur = ColumnsCursor({"brand_x": ["id"]})
|
||||
|
||||
assert vector_store._product_columns(cur, "brand_new") == "*"
|
||||
|
||||
|
||||
def test_the_column_probe_is_cached_until_invalidated():
|
||||
cur = ColumnsCursor({"brand_x": ["id", "embedding"]})
|
||||
|
||||
vector_store._product_columns(cur, "brand_x")
|
||||
vector_store._product_columns(cur, "brand_x")
|
||||
probes = [s for s in cur.statements if "information_schema.columns" in s]
|
||||
assert len(probes) == 1
|
||||
|
||||
vector_store._invalidate_product_columns_cache()
|
||||
vector_store._product_columns(cur, "brand_x")
|
||||
probes = [s for s in cur.statements if "information_schema.columns" in s]
|
||||
assert len(probes) == 2
|
||||
|
||||
|
||||
class _Conn:
|
||||
def __init__(self, cur):
|
||||
self._cur = cur
|
||||
|
||||
def cursor(self):
|
||||
return self
|
||||
|
||||
def __enter__(self):
|
||||
return self._cur
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_get_products_by_brand_selects_by_name_and_never_the_vectors(monkeypatch):
|
||||
cur = ColumnsCursor({"brand_cadbury": ["id", "product_name", "image_url", "embedding", "img_vector", "img_vector_src"]})
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _Conn(cur))
|
||||
|
||||
vector_store.get_products_by_brand("Cadbury")
|
||||
|
||||
select = [s for s in cur.statements if s.startswith("SELECT") and "FROM brand_cadbury" in s]
|
||||
assert select, cur.statements
|
||||
assert select[0].startswith('SELECT "id", "product_name", "image_url", "img_vector_src" FROM brand_cadbury')
|
||||
assert "embedding" not in select[0]
|
||||
assert '"img_vector"' not in select[0]
|
||||
|
||||
|
||||
def test_get_product_by_image_id_uses_the_same_projection(monkeypatch):
|
||||
cur = ColumnsCursor({"brand_cadbury": ["id", "embedding", "img_vector"]})
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _Conn(cur))
|
||||
|
||||
vector_store.get_product_by_image_id("Cadbury", "cadbury_x")
|
||||
|
||||
select = [s for s in cur.statements if "WHERE image_id = %s LIMIT 1" in s]
|
||||
assert select and select[0].startswith('SELECT "id" FROM brand_cadbury')
|
||||
|
||||
|
||||
def test_no_product_reader_is_select_star_any_more():
|
||||
import inspect
|
||||
for fn in (vector_store.get_products_by_brand, vector_store.get_product_by_image_id,
|
||||
vector_store.semantic_search, vector_store.text_search, vector_store.lexical_search):
|
||||
assert "SELECT *" not in inspect.getsource(fn), fn.__name__
|
||||
|
||||
|
||||
def test_the_mcp_slimmer_hides_the_vector_columns():
|
||||
from app.mcp_server import _slim
|
||||
|
||||
out = _slim({"title": "x", "img_vector": "[1,2]", "img_vector_src": "https://a/x.jpg", "embedding": "[0.1]"})
|
||||
|
||||
assert out == {"title": "x"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. the hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_schedule_is_a_no_op_when_disabled(monkeypatch):
|
||||
monkeypatch.setattr(iv, "ENABLE_IMAGE_VECTORS", False)
|
||||
q: "queue.Queue" = queue.Queue(maxsize=4)
|
||||
monkeypatch.setattr(iv, "_queue", q)
|
||||
started = []
|
||||
monkeypatch.setattr(iv, "_ensure_worker", lambda: started.append(1))
|
||||
|
||||
assert iv.schedule("Amul", ["a", "b"]) is False
|
||||
assert q.qsize() == 0 and not started
|
||||
|
||||
|
||||
def test_schedule_enqueues_and_the_worker_calls_backfill(monkeypatch):
|
||||
monkeypatch.setattr(iv, "ENABLE_IMAGE_VECTORS", True)
|
||||
q: "queue.Queue" = queue.Queue(maxsize=4)
|
||||
monkeypatch.setattr(iv, "_queue", q)
|
||||
monkeypatch.setattr(iv, "_worker", None)
|
||||
done = threading.Event()
|
||||
calls: List[Any] = []
|
||||
|
||||
def fake_backfill(brand, ids, **kw):
|
||||
calls.append((brand, ids, kw))
|
||||
done.set()
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(iv, "backfill_rows", fake_backfill)
|
||||
|
||||
assert iv.schedule("Amul", ["a", "", "b"]) is True
|
||||
assert done.wait(5), "worker never ran"
|
||||
assert calls == [("Amul", ["a", "b"], {"recompute_stale": True})]
|
||||
assert q.unfinished_tasks == 0 or q.join() is None
|
||||
|
||||
|
||||
def test_a_full_queue_is_logged_and_dropped_not_raised(monkeypatch, caplog):
|
||||
monkeypatch.setattr(iv, "ENABLE_IMAGE_VECTORS", True)
|
||||
q: "queue.Queue" = queue.Queue(maxsize=1)
|
||||
q.put_nowait(("X", ["1"]))
|
||||
monkeypatch.setattr(iv, "_queue", q)
|
||||
monkeypatch.setattr(iv, "_ensure_worker", lambda: None)
|
||||
|
||||
with caplog.at_level("WARNING", logger=iv.__name__):
|
||||
assert iv.schedule("Amul", ["a"]) is False
|
||||
assert "queue full" in caplog.text
|
||||
|
||||
|
||||
def test_an_empty_id_list_is_not_queued(monkeypatch):
|
||||
monkeypatch.setattr(iv, "ENABLE_IMAGE_VECTORS", True)
|
||||
q: "queue.Queue" = queue.Queue(maxsize=4)
|
||||
monkeypatch.setattr(iv, "_queue", q)
|
||||
monkeypatch.setattr(iv, "_ensure_worker", lambda: None)
|
||||
|
||||
assert iv.schedule("Amul", ["", None]) is False
|
||||
assert q.qsize() == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. the backfill: which rows, and dry run writes nothing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class RowsCursor:
|
||||
"""Returns canned rows for the candidate SELECT and records UPDATEs."""
|
||||
|
||||
def __init__(self, rows: List[Dict[str, Any]], has_columns: bool = True):
|
||||
self._rows = rows
|
||||
self._has = has_columns
|
||||
self.statements: List[tuple] = []
|
||||
self.description = None
|
||||
self._pending: Any = None
|
||||
|
||||
def execute(self, sql, params=None):
|
||||
text = " ".join(str(sql).split())
|
||||
self.statements.append((text, params))
|
||||
if "information_schema.columns" in text:
|
||||
self._pending = [("img_vector",), ("img_vector_src",)] if self._has else []
|
||||
elif text.startswith("SELECT image_id"):
|
||||
keys = ["image_id", "product_name", "image_url", "image_urls", "img_vector_src", "missing"]
|
||||
self.description = [(k,) for k in keys]
|
||||
self._pending = [tuple(r.get(k) for k in keys) for r in self._rows]
|
||||
else:
|
||||
self._pending = []
|
||||
|
||||
def fetchall(self):
|
||||
return self._pending or []
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
|
||||
_ROWS = [
|
||||
{"image_id": "missing", "image_url": "https://a/1.jpg", "image_urls": [], "img_vector_src": None, "missing": True},
|
||||
{"image_id": "stale", "image_url": "https://a/2-new.jpg", "image_urls": [], "img_vector_src": "https://a/2-old.jpg", "missing": False},
|
||||
{"image_id": "fresh", "image_url": "https://a/3.jpg", "image_urls": [], "img_vector_src": "https://a/3.jpg", "missing": False},
|
||||
{"image_id": "gallery-only", "image_url": "", "image_urls": ["https://a/4.jpg"], "img_vector_src": None, "missing": True},
|
||||
{"image_id": "no-image", "image_url": None, "image_urls": [], "img_vector_src": None, "missing": True},
|
||||
]
|
||||
|
||||
|
||||
def test_rows_needing_vectors_picks_missing_and_stale_with_a_primary():
|
||||
cur = RowsCursor(_ROWS)
|
||||
|
||||
got = iv.rows_needing_vectors(cur, "brand_a", recompute_stale=True)
|
||||
|
||||
assert [r["image_id"] for r in got] == ["missing", "stale", "gallery-only"]
|
||||
assert got[2]["primary_url"] == "https://a/4.jpg"
|
||||
|
||||
|
||||
def test_rows_needing_vectors_leaves_stale_rows_alone_unless_asked():
|
||||
cur = RowsCursor(_ROWS)
|
||||
|
||||
got = iv.rows_needing_vectors(cur, "brand_a", recompute_stale=False)
|
||||
|
||||
assert [r["image_id"] for r in got] == ["missing", "gallery-only"]
|
||||
|
||||
|
||||
def test_rows_needing_vectors_scopes_to_the_ids_just_written():
|
||||
cur = RowsCursor(_ROWS)
|
||||
|
||||
iv.rows_needing_vectors(cur, "brand_a", ["missing", "stale"])
|
||||
|
||||
text, params = cur.statements[-1]
|
||||
assert "image_id = ANY(%s)" in text and params == [["missing", "stale"]]
|
||||
|
||||
|
||||
def test_store_vector_does_not_touch_updated_at():
|
||||
cur = RowsCursor([])
|
||||
|
||||
iv.store_vector(cur, "brand_a", "x", [1, 2, 3], "https://a/x.jpg")
|
||||
|
||||
text, params = cur.statements[-1]
|
||||
assert text == "UPDATE brand_a SET img_vector = %s::vector, img_vector_src = %s WHERE image_id = %s"
|
||||
assert params == ("[1,2,3]", "https://a/x.jpg", "x")
|
||||
assert "updated_at" not in text
|
||||
|
||||
|
||||
class _RowsConn:
|
||||
def __init__(self, cur):
|
||||
self.cur = cur
|
||||
|
||||
def cursor(self):
|
||||
return self.cur
|
||||
|
||||
def close(self):
|
||||
pass
|
||||
|
||||
|
||||
def test_backfill_dry_run_downloads_and_writes_nothing(monkeypatch):
|
||||
cur = RowsCursor(_ROWS)
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _RowsConn(cur))
|
||||
fetched: List[str] = []
|
||||
monkeypatch.setattr(iv, "vector_for_url", lambda url, timeout=None: (fetched.append(url) or [0] * 3072, url))
|
||||
|
||||
totals = iv.backfill_rows("A", dry_run=True, pause_seconds=0)
|
||||
|
||||
assert totals["candidates"] == 3 and totals["computed"] == 3 and totals["failed"] == 0
|
||||
assert not any(t.startswith("UPDATE") for t, _ in cur.statements)
|
||||
|
||||
|
||||
def test_backfill_apply_writes_one_update_per_success_and_counts_failures(monkeypatch):
|
||||
cur = RowsCursor(_ROWS)
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _RowsConn(cur))
|
||||
monkeypatch.setattr(
|
||||
iv, "vector_for_url",
|
||||
lambda url, timeout=None: (None, url) if "4.jpg" in url else ([7] * 3072, url),
|
||||
)
|
||||
|
||||
totals = iv.backfill_rows("A", pause_seconds=0)
|
||||
|
||||
updates = [(t, p) for t, p in cur.statements if t.startswith("UPDATE")]
|
||||
assert totals == {"candidates": 3, "computed": 2, "failed": 1, "skipped": 0}
|
||||
assert [p[2] for _, p in updates] == ["missing", "stale"]
|
||||
assert updates[1][1][1] == "https://a/2-new.jpg"
|
||||
|
||||
|
||||
def test_backfill_skips_a_table_the_migration_has_not_reached(monkeypatch):
|
||||
cur = RowsCursor(_ROWS, has_columns=False)
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _RowsConn(cur))
|
||||
|
||||
totals = iv.backfill_rows("A", pause_seconds=0)
|
||||
|
||||
assert totals["skipped"] == 1 and totals["candidates"] == 0
|
||||
assert not any(t.startswith("SELECT image_id") for t, _ in cur.statements)
|
||||
|
||||
|
||||
def test_backfill_never_raises(monkeypatch):
|
||||
def boom():
|
||||
raise RuntimeError("db exploded")
|
||||
|
||||
monkeypatch.setattr(vector_store, "_connect", lambda: _RowsConn(type("C", (), {"execute": lambda *a: boom()})()))
|
||||
|
||||
assert iv.backfill_rows("A", pause_seconds=0)["computed"] == 0
|
||||
Reference in New Issue
Block a user