Vito Sansevero
a20e59f674
A Skipped tags list in System > Auto-tagging (D202) hides every suggested tag whose words contain an entry's words, on every suggestion path (Tag image, Auto-tag selected, the System batch, the legacy single and batch routes), before the limit. Jobs snapshot the list at their start. A read-only preview route counts what a list hides among the library's tags and the downloaded WD14 label lists. The draft store compares list and object drafts by value (sameDraft).
602 lines
24 KiB
Python
602 lines
24 KiB
Python
"""Scored top-N tag suggestions for a captured image (#249 B2).
|
||
|
||
``POST /prompt_manager/image-annotations/suggest`` runs the selected legacy
|
||
engine on the private capture (never the current file), behind the one
|
||
process-wide model gate shared with the legacy Auto-tag routes. It never
|
||
downloads a model and never writes. No real model is loaded: the engine is a
|
||
fake service.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import math
|
||
import os
|
||
import threading
|
||
import unittest
|
||
import uuid
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
try:
|
||
from test_api_annotations import AnnotationAPITestCase
|
||
from test_api_images import make_png
|
||
except ImportError:
|
||
from tests.test_api_annotations import AnnotationAPITestCase
|
||
from tests.test_api_images import make_png
|
||
|
||
try: # guarded so a missing module fails assertions, not test loading
|
||
from py import model_gate
|
||
except ImportError:
|
||
model_gate = None
|
||
|
||
from py import config as config_module
|
||
from py.autotag import MissingBackendPackages, package_fields
|
||
|
||
AutotagConfig = getattr(config_module, "AutotagConfig", None)
|
||
|
||
SUGGEST = "/prompt_manager/image-annotations/suggest"
|
||
WAIT = 30.0
|
||
WD14 = {"general_threshold": 0.35, "character_threshold": 0.85}
|
||
LLM = {"instructions": "Generate tags."}
|
||
UNAVAILABLE = "Download it in System › Auto-tagging."
|
||
NEEDS_LLAMA = (
|
||
"GGUF JoyCaption needs llama-cpp-python. Install it into ComfyUI's Python: "
|
||
'pip install llama-cpp-python. For a CUDA build: CMAKE_ARGS="-DGGML_CUDA=on" '
|
||
"pip install llama-cpp-python"
|
||
)
|
||
|
||
SCORED = [("sandstone", 0.98), ("desert", 0.96)] + [
|
||
(f"tag_{i:03d}", round(0.9 - i * 0.005, 4)) for i in range(150)
|
||
]
|
||
|
||
|
||
class FakeTagger:
|
||
"""Stand-in for ``AutoTagService`` with the suggestion entry point."""
|
||
|
||
default_prompt = "server default prompt"
|
||
custom_prompt = "legacy prompt"
|
||
wd14_general_threshold = 0.35
|
||
wd14_character_threshold = 0.85
|
||
|
||
def __init__(self):
|
||
self.models_config = {
|
||
"wd14-swinv2": {},
|
||
"wd14-vit": {},
|
||
"gguf": {},
|
||
"hf": {},
|
||
}
|
||
self.downloaded = {name: True for name in self.models_config}
|
||
# pip names of the backend packages that cannot be imported (#424)
|
||
self.missing = {name: [] for name in self.models_config}
|
||
self.loaded = None
|
||
self.load_calls = []
|
||
self.unload_calls = 0
|
||
self.calls = []
|
||
self.read_paths = []
|
||
self.read_bytes = []
|
||
self.skips = []
|
||
self.download_model = MagicMock(side_effect=AssertionError("downloaded"))
|
||
self.load_error = None
|
||
self.block = None # (entered, proceed) events
|
||
|
||
def get_models_status(self):
|
||
return {
|
||
name: {
|
||
"downloaded": self.downloaded[name],
|
||
"model_path": None,
|
||
**package_fields(name, self.missing[name]),
|
||
}
|
||
for name in self.models_config
|
||
}
|
||
|
||
def is_model_loaded(self):
|
||
return self.loaded is not None
|
||
|
||
def get_loaded_model_type(self):
|
||
return self.loaded
|
||
|
||
def load_model(self, model_type, use_gpu=True):
|
||
self.load_calls.append((model_type, use_gpu))
|
||
if self.load_error is not None:
|
||
raise self.load_error
|
||
self.loaded = model_type
|
||
return True
|
||
|
||
def unload_model(self):
|
||
self.unload_calls += 1
|
||
self.loaded = None
|
||
|
||
def generate_tags(self, image_path, **kwargs):
|
||
return ["legacy"]
|
||
|
||
def generate_tag_suggestions(
|
||
self,
|
||
image_path,
|
||
*,
|
||
limit,
|
||
prompt=None,
|
||
general_threshold=None,
|
||
character_threshold=None,
|
||
skip=None,
|
||
):
|
||
self.read_paths.append(image_path)
|
||
self.skips.append(skip)
|
||
with open(image_path, "rb") as handle:
|
||
self.read_bytes.append(handle.read())
|
||
self.calls.append(
|
||
{
|
||
"limit": limit,
|
||
"prompt": prompt,
|
||
"general_threshold": general_threshold,
|
||
"character_threshold": character_threshold,
|
||
"model": self.loaded,
|
||
}
|
||
)
|
||
if self.block is not None:
|
||
entered, proceed = self.block
|
||
entered.set()
|
||
proceed.wait(WAIT)
|
||
# As the engine does (#431): a skipped tag uses up no slot.
|
||
kept = [pair for pair in SCORED if skip is None or not skip.matches(pair[0])]
|
||
if self.loaded.startswith("wd14"):
|
||
return list(kept[:limit])
|
||
return [(name, None) for name, _ in kept[:limit]]
|
||
|
||
|
||
class SuggestTestCase(AnnotationAPITestCase):
|
||
async def get_application(self):
|
||
app = await super().get_application()
|
||
self.assertIsNotNone(model_gate, "py.model_gate is not implemented")
|
||
self.assertIsNotNone(AutotagConfig, "AutotagConfig is not implemented")
|
||
self.tagger = FakeTagger()
|
||
patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger)
|
||
patcher.start()
|
||
self.addCleanup(patcher.stop)
|
||
AutotagConfig.SUGGESTION_LIMIT = 10
|
||
self.addCleanup(setattr, AutotagConfig, "SUGGESTION_LIMIT", 10)
|
||
return app
|
||
|
||
async def tearDownAsync(self):
|
||
leaked = model_gate.holder() if model_gate is not None else None
|
||
if leaked is not None:
|
||
model_gate.release(leaked)
|
||
await super().tearDownAsync()
|
||
self.assertIsNone(leaked, "the suggest route left the model gate held")
|
||
|
||
async def observed(self, path="sub/a.png", size=(8, 6)):
|
||
make_png(self.output_dir / path, size=size)
|
||
return await self.observe_ok(path)
|
||
|
||
def suggest_body(self, observed, model_type="wd14-vit", options=None, **changes):
|
||
if options is None:
|
||
options = dict(WD14 if model_type.startswith("wd14") else LLM)
|
||
return {
|
||
"target": dict(observed["target"]),
|
||
"capture_token": observed["capture"]["token"],
|
||
"expected_epoch": observed["library_epoch"],
|
||
"model_type": model_type,
|
||
"options": options,
|
||
**changes,
|
||
}
|
||
|
||
async def suggest(self, body):
|
||
response = await self.client.post(SUGGEST, json=body)
|
||
return response.status, await self.json_of(response)
|
||
|
||
async def suggest_ok(self, body):
|
||
status, data = await self.suggest(body)
|
||
self.assertEqual(status, 200, data)
|
||
return data
|
||
|
||
def capture_path(self, observed):
|
||
return self.api._captures.get(observed["capture"]["token"]).path
|
||
|
||
async def wait_until_free(self):
|
||
for _ in range(3000):
|
||
if model_gate.holder() is None:
|
||
return
|
||
await asyncio.sleep(0.01)
|
||
self.fail(f"model gate still held by {model_gate.holder()}")
|
||
|
||
|
||
class TestSuggestions(SuggestTestCase):
|
||
async def test_wd14_suggestions_are_scored_top_n_from_the_capture(self):
|
||
observed = await self.observed()
|
||
|
||
body = await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
self.assertEqual(
|
||
body["suggestions"][:2],
|
||
[
|
||
{"name": "sandstone", "confidence": 0.98},
|
||
{"name": "desert", "confidence": 0.96},
|
||
],
|
||
)
|
||
self.assertEqual(len(body["suggestions"]), 10)
|
||
self.assertEqual(body["model_type"], "wd14-vit")
|
||
self.assertEqual(
|
||
body["options"],
|
||
{
|
||
"general_threshold": 0.35,
|
||
"character_threshold": 0.85,
|
||
"use_gpu": True,
|
||
"suggestion_limit": 10,
|
||
},
|
||
)
|
||
self.assertEqual(self.tagger.read_paths, [self.capture_path(observed)])
|
||
self.assertEqual(
|
||
self.tagger.calls[0],
|
||
{
|
||
"limit": 10,
|
||
"prompt": None,
|
||
"general_threshold": 0.35,
|
||
"character_threshold": 0.85,
|
||
"model": "wd14-vit",
|
||
},
|
||
)
|
||
self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)])
|
||
self.tagger.download_model.assert_not_called()
|
||
self.assertIsNone(model_gate.holder())
|
||
self.assert_no_paths(body)
|
||
|
||
async def test_suggestions_keep_the_requested_limit(self):
|
||
observed = await self.observed()
|
||
for limit in (1, 10, 100):
|
||
with self.subTest(limit=limit):
|
||
options = {**WD14, "suggestion_limit": limit}
|
||
body = await self.suggest_ok(
|
||
self.suggest_body(observed, options=options)
|
||
)
|
||
self.assertEqual(len(body["suggestions"]), limit)
|
||
self.assertEqual(body["options"]["suggestion_limit"], limit)
|
||
# Omitted: the persisted autotag.defaults.suggestion_limit.
|
||
AutotagConfig.SUGGESTION_LIMIT = 3
|
||
body = await self.suggest_ok(self.suggest_body(observed))
|
||
self.assertEqual(len(body["suggestions"]), 3)
|
||
self.assertEqual(body["options"]["suggestion_limit"], 3)
|
||
|
||
async def test_llm_suggestions_send_instructions_and_null_confidence(self):
|
||
observed = await self.observed()
|
||
self.tagger.custom_prompt = "legacy prompt"
|
||
for model_type in ("gguf", "hf"):
|
||
with self.subTest(model_type=model_type):
|
||
options = {"instructions": f"tags for {model_type}", "use_gpu": False}
|
||
body = await self.suggest_ok(
|
||
self.suggest_body(observed, model_type, options)
|
||
)
|
||
self.assertEqual(body["model_type"], model_type)
|
||
self.assertEqual(body["options"], {**options, "suggestion_limit": 10})
|
||
self.assertEqual(
|
||
body["suggestions"][:2],
|
||
[
|
||
{"name": "sandstone", "confidence": None},
|
||
{"name": "desert", "confidence": None},
|
||
],
|
||
)
|
||
call = self.tagger.calls[-1]
|
||
self.assertEqual(call["prompt"], f"tags for {model_type}")
|
||
self.assertIsNone(call["general_threshold"])
|
||
self.assertEqual(self.tagger.load_calls[-1], (model_type, False))
|
||
self.assertEqual(self.tagger.custom_prompt, "legacy prompt")
|
||
|
||
async def test_loaded_model_is_reused_and_another_is_swapped_in(self):
|
||
observed = await self.observed()
|
||
self.tagger.loaded = "wd14-vit"
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
await self.suggest_ok(self.suggest_body(observed, "wd14-swinv2"))
|
||
self.assertEqual(self.tagger.load_calls, [("wd14-swinv2", True)])
|
||
self.assertEqual(self.tagger.calls[-1]["model"], "wd14-swinv2")
|
||
|
||
async def test_suggest_uses_capture_bytes_not_current_file(self):
|
||
observed = await self.observed()
|
||
original = (self.output_dir / "sub" / "a.png").read_bytes()
|
||
make_png(self.output_dir / "sub" / "a.png", size=(9, 9), color=(0, 0, 255))
|
||
self.assertNotEqual((self.output_dir / "sub" / "a.png").read_bytes(), original)
|
||
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
capture_path = self.capture_path(observed)
|
||
self.assertEqual(self.tagger.read_paths, [capture_path])
|
||
self.assertEqual(self.tagger.read_bytes, [original])
|
||
self.assertNotEqual(
|
||
os.path.normcase(capture_path),
|
||
os.path.normcase(str(self.output_dir / "sub" / "a.png")),
|
||
)
|
||
|
||
async def test_suggest_writes_nothing(self):
|
||
observed = await self.observed()
|
||
before = self.db_snapshot()
|
||
rows = self.bytes_rows()
|
||
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
await self.suggest_ok(self.suggest_body(observed, "gguf"))
|
||
|
||
self.assertEqual(self.db_snapshot(), before)
|
||
self.assertEqual(self.bytes_rows(), rows)
|
||
# The capture stays usable for Save after Generate.
|
||
self.assertIsNotNone(self.api._captures.get(observed["capture"]["token"]))
|
||
|
||
|
||
class TestModelAvailability(SuggestTestCase):
|
||
async def test_suggest_never_downloads_and_reports_model_unavailable(self):
|
||
observed = await self.observed()
|
||
self.tagger.downloaded["wd14-vit"] = False
|
||
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (409, "model-unavailable"))
|
||
self.assertIn(UNAVAILABLE, body["error"])
|
||
self.assertIn("WD14 ViT", body["error"])
|
||
self.tagger.download_model.assert_not_called()
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
self.assertEqual(self.tagger.calls, [])
|
||
self.assertIsNone(model_gate.holder())
|
||
|
||
async def test_missing_backend_package_is_refused_before_the_model(self):
|
||
observed = await self.observed()
|
||
self.tagger.missing["gguf"] = ["llama-cpp-python"]
|
||
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
|
||
self.assertEqual((status, body["code"]), (409, "model-unavailable"))
|
||
self.assertEqual(body["success"], False)
|
||
self.assertEqual(body["error"], NEEDS_LLAMA)
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
self.assertIsNone(model_gate.holder())
|
||
# Not downloaded still reads as before, whatever the packages.
|
||
self.tagger.downloaded["gguf"] = False
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
self.assertIn("is not downloaded", body["error"])
|
||
|
||
async def test_a_package_lost_after_the_check_is_the_same_409_without_a_traceback(
|
||
self,
|
||
):
|
||
observed = await self.observed()
|
||
self.tagger.load_error = MissingBackendPackages("gguf", ["llama-cpp-python"])
|
||
|
||
with self.assertNoLogs(self.api.logger, level="ERROR"):
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
|
||
self.assertEqual((status, body["code"]), (409, "model-unavailable"))
|
||
self.assertEqual(body["error"], NEEDS_LLAMA)
|
||
self.assertIsNone(model_gate.holder())
|
||
entry = self.api._captures._entries[observed["capture"]["token"]]
|
||
self.assertEqual(entry.uses, 0)
|
||
|
||
async def test_load_failure_is_a_generic_error_and_frees_everything(self):
|
||
observed = await self.observed()
|
||
self.tagger.load_error = RuntimeError(f"cannot load {self.output_dir}")
|
||
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (500, "suggest-failed"))
|
||
self.assert_no_paths(body)
|
||
self.assertIsNone(model_gate.holder())
|
||
entry = self.api._captures._entries[observed["capture"]["token"]]
|
||
self.assertEqual(entry.uses, 0)
|
||
|
||
async def test_gate_work_that_never_starts_releases_the_capture_pin(self):
|
||
observed = await self.observed()
|
||
loop = asyncio.get_running_loop()
|
||
real = loop.run_in_executor
|
||
|
||
def refuse_gate_work(executor, func, *args):
|
||
if getattr(func, "__qualname__", "").startswith("ModelGate.run"):
|
||
raise RuntimeError("cannot schedule new futures after shutdown")
|
||
return real(executor, func, *args)
|
||
|
||
with patch.object(loop, "run_in_executor", refuse_gate_work):
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
|
||
self.assertEqual((status, body["code"]), (500, "suggest-failed"))
|
||
self.assertIsNone(model_gate.holder())
|
||
entry = self.api._captures._entries[observed["capture"]["token"]]
|
||
self.assertEqual(entry.uses, 0)
|
||
self.assertEqual(self.tagger.calls, [])
|
||
|
||
|
||
class TestModelGate(SuggestTestCase):
|
||
async def _blocked_suggest(self, observed, model_type="wd14-vit"):
|
||
entered, proceed = threading.Event(), threading.Event()
|
||
self.tagger.block = (entered, proceed)
|
||
task = asyncio.ensure_future(
|
||
self.client.post(SUGGEST, json=self.suggest_body(observed, model_type))
|
||
)
|
||
await asyncio.get_running_loop().run_in_executor(None, entered.wait, WAIT)
|
||
self.assertTrue(entered.is_set())
|
||
return task, proceed
|
||
|
||
async def test_suggest_reports_busy_while_legacy_batch_runs(self):
|
||
observed = await self.observed()
|
||
prompt_id = self._save_prompt("batch prompt")
|
||
self.api.db.link_image_to_prompt(prompt_id, str(self.output_dir / "sub/a.png"))
|
||
started, proceed = threading.Event(), threading.Event()
|
||
# The legacy batch is the autotag job (#379): it suggests per image.
|
||
self.tagger.block = (started, proceed)
|
||
batch = asyncio.ensure_future(
|
||
self.client.post("/prompt_manager/autotag/start?model_type=wd14-vit")
|
||
)
|
||
try:
|
||
await asyncio.get_running_loop().run_in_executor(None, started.wait, WAIT)
|
||
calls = len(self.tagger.calls)
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||
self.assertEqual(len(self.tagger.calls), calls)
|
||
self.assertEqual(self.tagger.loaded, "wd14-vit") # not swapped
|
||
finally:
|
||
self.tagger.block = None
|
||
proceed.set()
|
||
response = await batch
|
||
await response.read()
|
||
await self.wait_until_free()
|
||
await self.suggest_ok(self.suggest_body(observed))
|
||
|
||
async def test_legacy_unload_returns_model_busy_while_native_generate_runs(self):
|
||
observed = await self.observed()
|
||
task, proceed = await self._blocked_suggest(observed)
|
||
try:
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 409)
|
||
self.assertEqual((await response.json())["code"], "model-busy")
|
||
response = await self.client.post(
|
||
"/prompt_manager/autotag/single",
|
||
json={"root_index": 0, "path": "sub/a.png", "model_type": "gguf"},
|
||
)
|
||
self.assertEqual(response.status, 409)
|
||
status, body = await self.suggest(self.suggest_body(observed, "gguf"))
|
||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||
self.assertEqual(self.tagger.unload_calls, 0)
|
||
self.assertEqual(self.tagger.loaded, "wd14-vit")
|
||
finally:
|
||
proceed.set()
|
||
response = await task
|
||
self.assertEqual(response.status, 200)
|
||
await self.wait_until_free()
|
||
self.assertEqual([call["model"] for call in self.tagger.calls], ["wd14-vit"])
|
||
|
||
async def test_gate_held_after_client_disconnect_until_executor_finishes(self):
|
||
observed = await self.observed()
|
||
token = observed["capture"]["token"]
|
||
task, proceed = await self._blocked_suggest(observed)
|
||
task.cancel()
|
||
with self.assertRaises(asyncio.CancelledError):
|
||
await task
|
||
try:
|
||
# The client is gone; the inference still runs, so the gate and
|
||
# the capture pin stay until the worker thread exits.
|
||
self.assertIsNotNone(model_gate.holder())
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 409)
|
||
self.assertEqual(self.api._captures._entries[token].uses, 1)
|
||
finally:
|
||
proceed.set()
|
||
await self.wait_until_free()
|
||
for _ in range(3000):
|
||
if self.api._captures._entries[token].uses == 0:
|
||
break
|
||
await asyncio.sleep(0.01)
|
||
self.assertEqual(self.api._captures._entries[token].uses, 0)
|
||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||
self.assertEqual(response.status, 200)
|
||
|
||
|
||
class TestRequestValidation(SuggestTestCase):
|
||
async def test_strict_option_validation(self):
|
||
observed = await self.observed()
|
||
bad_options = {
|
||
"wd14-vit": [
|
||
{**WD14, "general_threshold": True},
|
||
{**WD14, "general_threshold": math.nan},
|
||
{**WD14, "general_threshold": math.inf},
|
||
{**WD14, "general_threshold": "0.3"},
|
||
{**WD14, "general_threshold": 1.5},
|
||
{**WD14, "character_threshold": -0.1},
|
||
{"general_threshold": 0.35},
|
||
{**WD14, "instructions": "x"},
|
||
{**WD14, "unknown": 1},
|
||
{**WD14, "use_gpu": "yes"},
|
||
{**WD14, "suggestion_limit": 0},
|
||
{**WD14, "suggestion_limit": 101},
|
||
{**WD14, "suggestion_limit": 2.5},
|
||
{**WD14, "suggestion_limit": "10"},
|
||
{**WD14, "suggestion_limit": True},
|
||
],
|
||
"gguf": [
|
||
{**LLM, "general_threshold": 0.35},
|
||
{**LLM, "character_threshold": 0.85},
|
||
{"instructions": "x" * 4001},
|
||
{"instructions": "\ud800"}, # not encodable as UTF-8
|
||
{"instructions": " "},
|
||
{"instructions": 5},
|
||
{},
|
||
{**LLM, "suggestion_limit": 0},
|
||
],
|
||
}
|
||
for model_type, cases in bad_options.items():
|
||
for options in cases:
|
||
with self.subTest(model_type=model_type, options=repr(options)[:80]):
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, model_type, options)
|
||
)
|
||
self.assertEqual(status, 400, body)
|
||
for body in (
|
||
{**self.suggest_body(observed), "extra": 1},
|
||
{k: v for k, v in self.suggest_body(observed).items() if k != "options"},
|
||
self.suggest_body(observed, options=[]),
|
||
self.suggest_body(observed, model_type="ollama"),
|
||
self.suggest_body(observed, model_type=7, options=dict(WD14)),
|
||
self.suggest_body(observed, capture_token="../x"),
|
||
self.suggest_body(observed, expected_epoch="nope"),
|
||
self.suggest_body(observed, target={"root_index": 0}),
|
||
):
|
||
with self.subTest(body=json.dumps(body)[:120]):
|
||
status, _ = await self.suggest(body)
|
||
self.assertEqual(status, 400)
|
||
# The edges that are allowed.
|
||
await self.suggest_ok(
|
||
self.suggest_body(
|
||
observed,
|
||
options={
|
||
"general_threshold": 0,
|
||
"character_threshold": 1,
|
||
"suggestion_limit": 100,
|
||
},
|
||
)
|
||
)
|
||
await self.suggest_ok(
|
||
self.suggest_body(observed, "hf", options={"instructions": "x" * 4000})
|
||
)
|
||
self.assertEqual(len(self.tagger.calls), 2)
|
||
self.tagger.download_model.assert_not_called()
|
||
|
||
async def test_capture_token_must_match_target_and_epoch(self):
|
||
observed = await self.observed()
|
||
make_png(self.output_dir / "sub" / "b.png")
|
||
other = await self.observe_ok("sub/b.png")
|
||
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, capture_token="x" * 43)
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, capture_token=other["capture"]["token"])
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
status, body = await self.suggest(
|
||
self.suggest_body(observed, expected_epoch=str(uuid.uuid4()))
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "storage-changed"))
|
||
self.assertEqual(body["library_epoch"], self.epoch)
|
||
status, body = await self.suggest(
|
||
self.suggest_body(
|
||
observed, target={**observed["target"], "root_scope": "0" * 64}
|
||
)
|
||
)
|
||
self.assertEqual((status, body["code"]), (409, "source-changed"))
|
||
self.api._captures.release(observed["capture"]["token"])
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
self.assertEqual((status, body["code"]), (409, "capture-expired"))
|
||
self.assertEqual(self.tagger.calls, [])
|
||
for token in (observed, other):
|
||
entry = self.api._captures._entries.get(token["capture"]["token"])
|
||
self.assertTrue(entry is None or entry.uses == 0)
|
||
|
||
async def test_video_and_oversized_dimensions_unsupported_media(self):
|
||
(self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 4)
|
||
video = await self.observe_ok("clip.mp4")
|
||
self.assertFalse(video["capture"]["generate"]["supported"])
|
||
big = await self.observed("big.png", size=(4097, 1))
|
||
self.assertFalse(big["capture"]["generate"]["supported"])
|
||
|
||
for observed in (video, big):
|
||
status, body = await self.suggest(self.suggest_body(observed))
|
||
self.assertEqual((status, body["code"]), (415, "unsupported-media"))
|
||
self.assertEqual(body["error"], observed["capture"]["generate"]["reason"])
|
||
self.assertEqual(self.tagger.calls, [])
|
||
self.assertEqual(self.tagger.load_calls, [])
|
||
self.assertIsNone(model_gate.holder())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|