test: rename SearchSpace -> Workspace across tests + fixtures (Phase 2 Wave F)

Apply the same rename to surfsense_backend/tests: workspace_id fields/vars,
Workspace* classes/schemas, table name searchspaces -> workspaces in raw SQL,
and the API URL spellings -> /workspaces. Preserves the carve-out wire literals
tests assert (Celery task names, OTel key "search_space.id").
This commit is contained in:
CREDO23 2026-06-26 18:36:46 +02:00
parent 56826a63bc
commit ca9bd28934
124 changed files with 1269 additions and 1269 deletions

View file

@ -1,4 +1,4 @@
"""Unit tests for ``_resolve_agent_billing_for_search_space``."""
"""Unit tests for ``_resolve_agent_billing_for_workspace``."""
from __future__ import annotations
@ -39,7 +39,7 @@ class _FakePinResolution:
from_existing_pin: bool = False
def _make_search_space(*, chat_model_id: int | None, user_id: UUID) -> SimpleNamespace:
def _make_workspace(*, chat_model_id: int | None, user_id: UUID) -> SimpleNamespace:
return SimpleNamespace(id=42, chat_model_id=chat_model_id, user_id=user_id)
@ -50,16 +50,16 @@ def _make_byok_model(
id=id_,
model_id=model_id,
catalog={"base_model": base_model} if base_model else {},
connection=SimpleNamespace(enabled=True, search_space_id=42, user_id=None),
connection=SimpleNamespace(enabled=True, workspace_id=42, user_id=None),
)
@pytest.mark.asyncio
async def test_auto_mode_with_thread_id_resolves_to_premium_global(monkeypatch):
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=0, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=0, user_id=user_id)])
async def _fake_resolve_pin(*_args, **kwargs):
assert kwargs["selected_llm_config_id"] == 0
@ -84,8 +84,8 @@ async def test_auto_mode_with_thread_id_resolves_to_premium_global(monkeypatch):
)
monkeypatch.setattr(llm_module, "get_global_llm_config", _fake_get_global)
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42, thread_id=99
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42, thread_id=99
)
assert owner == user_id
@ -95,14 +95,14 @@ async def test_auto_mode_with_thread_id_resolves_to_premium_global(monkeypatch):
@pytest.mark.asyncio
async def test_auto_mode_with_thread_id_resolves_to_byok_is_free(monkeypatch):
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
search_space = _make_search_space(chat_model_id=0, user_id=user_id)
workspace = _make_workspace(chat_model_id=0, user_id=user_id)
byok_model = _make_byok_model(
id_=17, base_model="anthropic/claude-3-haiku", model_id="my-claude"
)
session = _FakeSession([search_space, byok_model])
session = _FakeSession([workspace, byok_model])
async def _fake_resolve_pin(*_args, **_kwargs):
return _FakePinResolution(resolved_llm_config_id=17, resolved_tier="free")
@ -113,8 +113,8 @@ async def test_auto_mode_with_thread_id_resolves_to_byok_is_free(monkeypatch):
pin_module, "resolve_or_get_pinned_llm_config_id", _fake_resolve_pin
)
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42, thread_id=99
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42, thread_id=99
)
assert owner == user_id
@ -124,13 +124,13 @@ async def test_auto_mode_with_thread_id_resolves_to_byok_is_free(monkeypatch):
@pytest.mark.asyncio
async def test_auto_mode_without_thread_id_falls_back_to_free():
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=0, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=0, user_id=user_id)])
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42, thread_id=None
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42, thread_id=None
)
assert owner == user_id
@ -140,10 +140,10 @@ async def test_auto_mode_without_thread_id_falls_back_to_free():
@pytest.mark.asyncio
async def test_auto_mode_pin_failure_falls_back_to_free(monkeypatch):
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=0, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=0, user_id=user_id)])
async def _fake_resolve_pin(*args, **kwargs):
raise ValueError("thread missing")
@ -154,8 +154,8 @@ async def test_auto_mode_pin_failure_falls_back_to_free(monkeypatch):
pin_module, "resolve_or_get_pinned_llm_config_id", _fake_resolve_pin
)
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42, thread_id=99
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42, thread_id=99
)
assert owner == user_id
@ -165,10 +165,10 @@ async def test_auto_mode_pin_failure_falls_back_to_free(monkeypatch):
@pytest.mark.asyncio
async def test_negative_id_premium_global_returns_premium(monkeypatch):
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=-1, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=-1, user_id=user_id)])
def _fake_get_global(cfg_id):
return {
@ -182,8 +182,8 @@ async def test_negative_id_premium_global_returns_premium(monkeypatch):
monkeypatch.setattr(llm_module, "get_global_llm_config", _fake_get_global)
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42, thread_id=99
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42, thread_id=99
)
assert owner == user_id
@ -193,10 +193,10 @@ async def test_negative_id_premium_global_returns_premium(monkeypatch):
@pytest.mark.asyncio
async def test_negative_id_missing_base_model_falls_back_to_model_name(monkeypatch):
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=-5, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=-5, user_id=user_id)])
def _fake_get_global(cfg_id):
return {"id": cfg_id, "model_name": "fallback-model", "billing_tier": "premium"}
@ -205,8 +205,8 @@ async def test_negative_id_missing_base_model_falls_back_to_model_name(monkeypat
monkeypatch.setattr(llm_module, "get_global_llm_config", _fake_get_global)
_, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42
_, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42
)
assert tier == "premium"
@ -215,15 +215,15 @@ async def test_negative_id_missing_base_model_falls_back_to_model_name(monkeypat
@pytest.mark.asyncio
async def test_positive_id_byok_is_always_free():
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
search_space = _make_search_space(chat_model_id=23, user_id=user_id)
workspace = _make_workspace(chat_model_id=23, user_id=user_id)
byok_model = _make_byok_model(id_=23, base_model="anthropic/claude-3.5-sonnet")
session = _FakeSession([search_space, byok_model])
session = _FakeSession([workspace, byok_model])
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42
)
assert owner == user_id
@ -233,13 +233,13 @@ async def test_positive_id_byok_is_always_free():
@pytest.mark.asyncio
async def test_positive_id_byok_missing_returns_free_with_empty_base_model():
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=99, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=99, user_id=user_id)])
owner, tier, base_model = await _resolve_agent_billing_for_search_space(
session, search_space_id=42
owner, tier, base_model = await _resolve_agent_billing_for_workspace(
session, workspace_id=42
)
assert owner == user_id
@ -248,21 +248,21 @@ async def test_positive_id_byok_missing_returns_free_with_empty_base_model():
@pytest.mark.asyncio
async def test_search_space_not_found_raises_value_error():
from app.services.billable_calls import _resolve_agent_billing_for_search_space
async def test_workspace_not_found_raises_value_error():
from app.services.billable_calls import _resolve_agent_billing_for_workspace
with pytest.raises(ValueError, match="Search space"):
await _resolve_agent_billing_for_search_space(
_FakeSession([None]), search_space_id=999
with pytest.raises(ValueError, match="Workspace"):
await _resolve_agent_billing_for_workspace(
_FakeSession([None]), workspace_id=999
)
@pytest.mark.asyncio
async def test_chat_model_id_none_raises_value_error():
from app.services.billable_calls import _resolve_agent_billing_for_search_space
from app.services.billable_calls import _resolve_agent_billing_for_workspace
user_id = uuid4()
session = _FakeSession([_make_search_space(chat_model_id=None, user_id=user_id)])
session = _FakeSession([_make_workspace(chat_model_id=None, user_id=user_id)])
with pytest.raises(ValueError, match="chat_model_id"):
await _resolve_agent_billing_for_search_space(session, search_space_id=42)
await _resolve_agent_billing_for_workspace(session, workspace_id=42)

View file

@ -11,7 +11,7 @@ def test_lock_key_format():
from app.tasks.celery_tasks.document_tasks import _ai_sort_lock_key
key = _ai_sort_lock_key(42)
assert key == "ai_sort:search_space:42:lock"
assert key == "ai_sort:workspace:42:lock"
def test_lock_prevents_duplicate_run():
@ -31,11 +31,11 @@ def test_lock_prevents_duplicate_run():
):
import asyncio
from app.tasks.celery_tasks.document_tasks import _ai_sort_search_space_async
from app.tasks.celery_tasks.document_tasks import _ai_sort_workspace_async
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(_ai_sort_search_space_async(1, "user-123"))
loop.run_until_complete(_ai_sort_workspace_async(1, "user-123"))
finally:
loop.close()

View file

@ -140,12 +140,12 @@ def _set_global_llm_configs(monkeypatch, config, configs: list[dict]):
def _thread(
*,
search_space_id: int = 10,
workspace_id: int = 10,
pinned_llm_config_id: int | None = None,
):
return SimpleNamespace(
id=1,
search_space_id=search_space_id,
workspace_id=workspace_id,
pinned_llm_config_id=pinned_llm_config_id,
)
@ -186,7 +186,7 @@ async def test_auto_first_turn_pins_one_model(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -234,7 +234,7 @@ async def test_premium_eligible_auto_prefers_premium_over_free(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -321,7 +321,7 @@ async def test_premium_eligible_auto_uses_quality_pool_not_single_preferred_mode
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -361,7 +361,7 @@ async def test_next_turn_reuses_existing_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -400,7 +400,7 @@ async def test_premium_eligible_auto_can_pin_premium(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -445,7 +445,7 @@ async def test_premium_ineligible_auto_pins_free_only(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -490,7 +490,7 @@ async def test_pinned_premium_stays_premium_after_quota_exhaustion(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -535,7 +535,7 @@ async def test_force_repin_free_switches_auto_premium_pin_to_free(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
force_repin_free=True,
@ -567,7 +567,7 @@ async def test_explicit_user_model_change_clears_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=7,
)
@ -605,7 +605,7 @@ async def test_invalid_pinned_config_repairs_with_new_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -664,7 +664,7 @@ async def test_health_gated_config_is_excluded_from_selection(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -716,7 +716,7 @@ async def test_tier_a_locks_first_premium_user_skips_or(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -768,7 +768,7 @@ async def test_tier_a_falls_through_to_or_when_a_pool_empty_for_user(monkeypatch
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -825,7 +825,7 @@ async def test_top_k_picks_only_high_score_models(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=thread_id,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -889,7 +889,7 @@ async def test_pin_reuse_survives_health_gating_for_existing_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -941,7 +941,7 @@ async def test_pin_reuse_regression_existing_healthy_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -1001,7 +1001,7 @@ async def test_runtime_cooled_down_pin_is_not_reused(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -1069,7 +1069,7 @@ async def test_shared_runtime_cooldown_blocks_pin_across_workers(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -1113,7 +1113,7 @@ async def test_clearing_runtime_cooldown_restores_pin_reuse(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
)
@ -1165,7 +1165,7 @@ async def test_auto_pin_repin_excludes_previous_config_on_runtime_retry(monkeypa
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id="00000000-0000-0000-0000-000000000001",
selected_llm_config_id=0,
exclude_config_ids={-1},

View file

@ -76,7 +76,7 @@ class _FakeSession:
def _thread(*, pinned: int | None = None):
return SimpleNamespace(id=1, search_space_id=10, pinned_llm_config_id=pinned)
return SimpleNamespace(id=1, workspace_id=10, pinned_llm_config_id=pinned)
def _set_global_llm_configs(monkeypatch, config, configs: list[dict]):
@ -179,7 +179,7 @@ async def test_image_turn_filters_out_text_only_candidates(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
requires_image_input=True,
@ -207,7 +207,7 @@ async def test_image_turn_force_repins_stale_text_only_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
requires_image_input=True,
@ -239,7 +239,7 @@ async def test_image_turn_reuses_existing_vision_pin(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
requires_image_input=True,
@ -269,7 +269,7 @@ async def test_image_turn_with_no_vision_candidates_raises(monkeypatch):
await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
requires_image_input=True,
@ -292,7 +292,7 @@ async def test_non_image_turn_keeps_text_only_in_pool(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
)
@ -326,7 +326,7 @@ async def test_image_turn_unannotated_cfg_resolves_via_helper(monkeypatch):
result = await resolve_or_get_pinned_llm_config_id(
session,
thread_id=1,
search_space_id=10,
workspace_id=10,
user_id=None,
selected_llm_config_id=0,
requires_image_input=True,

View file

@ -169,7 +169,7 @@ async def test_free_path_skips_reserve_but_writes_audit_row(monkeypatch):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="free",
base_model="openai/gpt-image-1",
usage_type="image_generation",
@ -210,7 +210,7 @@ async def test_premium_reserve_denied_raises_quota_insufficient(monkeypatch):
with pytest.raises(QuotaInsufficientError) as exc_info:
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="openai/gpt-image-1",
quota_reserve_micros_override=50_000,
@ -242,7 +242,7 @@ async def test_premium_success_finalizes_with_actual_cost(monkeypatch):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="openai/gpt-image-1",
quota_reserve_micros_override=50_000,
@ -292,7 +292,7 @@ async def test_premium_failure_releases_reservation(monkeypatch):
with pytest.raises(_ProviderError):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="openai/gpt-image-1",
quota_reserve_micros_override=50_000,
@ -336,7 +336,7 @@ async def test_premium_uses_estimator_when_no_micros_override(monkeypatch):
user_id = uuid4()
async with billable_call(
user_id=user_id,
search_space_id=1,
workspace_id=1,
billing_tier="premium",
base_model="openai/gpt-4o",
quota_reserve_tokens=4000,
@ -367,7 +367,7 @@ async def test_premium_finalize_failure_propagates_and_releases(monkeypatch):
with pytest.raises(BillingSettlementError):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="openai/gpt-image-1",
quota_reserve_micros_override=50_000,
@ -408,7 +408,7 @@ async def test_premium_audit_commit_hang_times_out_after_finalize(monkeypatch):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="openai/gpt-image-1",
quota_reserve_micros_override=50_000,
@ -450,7 +450,7 @@ async def test_free_audit_failure_is_best_effort(monkeypatch):
async with billable_call(
user_id=uuid4(),
search_space_id=42,
workspace_id=42,
billing_tier="free",
base_model="openai/gpt-image-1",
usage_type="image_generation",
@ -488,7 +488,7 @@ async def test_free_podcast_path_audits_with_podcast_usage_type(monkeypatch):
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="free",
base_model="openrouter/some-free-model",
quota_reserve_micros_override=200_000,
@ -514,7 +514,7 @@ async def test_free_podcast_path_audits_with_podcast_usage_type(monkeypatch):
row = spies["record"][0]
assert row["usage_type"] == "podcast_generation"
assert row["thread_id"] is None
assert row["search_space_id"] == 42
assert row["workspace_id"] == 42
assert row["call_details"] == {"podcast_id": 7, "title": "Test Podcast"}
@ -539,7 +539,7 @@ async def test_premium_video_denial_raises_quota_insufficient(monkeypatch):
with pytest.raises(QuotaInsufficientError) as exc_info:
async with billable_call(
user_id=user_id,
search_space_id=42,
workspace_id=42,
billing_tier="premium",
base_model="gpt-5.4",
quota_reserve_micros_override=1_000_000,

View file

@ -46,8 +46,8 @@ async def test_global_openrouter_image_gen_sets_explicit_api_base():
image_gen.response_format = None
image_gen.model = None
search_space = MagicMock()
search_space.image_gen_model_id = global_model["id"]
workspace = MagicMock()
workspace.image_gen_model_id = global_model["id"]
session = MagicMock()
with (
@ -68,7 +68,7 @@ async def test_global_openrouter_image_gen_sets_explicit_api_base():
),
):
await image_generation_routes._execute_image_generation(
session=session, image_gen=image_gen, search_space=search_space
session=session, image_gen=image_gen, workspace=workspace
)
assert captured.get("api_base") == "https://openrouter.ai/api/v1"
@ -109,16 +109,16 @@ async def test_generate_image_tool_global_sets_explicit_api_base():
response._hidden_params = {"model": "openrouter/openai/gpt-image-1"}
return response
search_space = MagicMock()
search_space.id = 1
search_space.image_gen_model_id = global_model["id"]
workspace = MagicMock()
workspace.id = 1
workspace.image_gen_model_id = global_model["id"]
session_cm = AsyncMock()
session = AsyncMock()
session_cm.__aenter__.return_value = session
scalars = MagicMock()
scalars.first.return_value = search_space
scalars.first.return_value = workspace
exec_result = MagicMock()
exec_result.scalars.return_value = scalars
session.execute.return_value = exec_result
@ -146,7 +146,7 @@ async def test_generate_image_tool_global_sets_explicit_api_base():
),
):
tool = gi_module.create_generate_image_tool(
search_space_id=1, db_session=MagicMock()
workspace_id=1, db_session=MagicMock()
)
# The live tool takes an injected ToolRuntime and returns a Command;
# drive the raw coroutine with a minimal runtime (the tool only reads

View file

@ -77,7 +77,7 @@ async def test_ainvoke_routes_through_billable_call(monkeypatch):
wrapper = QuotaCheckedVisionLLM(
inner,
user_id=user_id,
search_space_id=99,
workspace_id=99,
billing_tier="premium",
base_model="openai/gpt-4o",
quota_reserve_tokens=4000,
@ -89,7 +89,7 @@ async def test_ainvoke_routes_through_billable_call(monkeypatch):
assert len(captured_kwargs) == 1
bc_kwargs = captured_kwargs[0]
assert bc_kwargs["user_id"] == user_id
assert bc_kwargs["search_space_id"] == 99
assert bc_kwargs["workspace_id"] == 99
assert bc_kwargs["billing_tier"] == "premium"
assert bc_kwargs["base_model"] == "openai/gpt-4o"
assert bc_kwargs["quota_reserve_tokens"] == 4000
@ -120,7 +120,7 @@ async def test_ainvoke_propagates_quota_insufficient_error(monkeypatch):
wrapper = QuotaCheckedVisionLLM(
inner,
user_id=uuid4(),
search_space_id=1,
workspace_id=1,
billing_tier="premium",
base_model="openai/gpt-4o",
quota_reserve_tokens=4000,
@ -146,7 +146,7 @@ async def test_proxies_non_overridden_attributes_to_inner():
wrapper = QuotaCheckedVisionLLM(
inner,
user_id=uuid4(),
search_space_id=1,
workspace_id=1,
billing_tier="premium",
base_model="openai/gpt-4o",
quota_reserve_tokens=4000,

View file

@ -93,7 +93,7 @@ def _action(*, tool_name: str, action_id: int = 7):
id=action_id,
tool_name=tool_name,
thread_id=1,
search_space_id=2,
workspace_id=2,
user_id="user-1",
reverse_descriptor=None,
)
@ -111,7 +111,7 @@ def _doc_revision(
revision = MagicMock()
revision.id = 100
revision.document_id = document_id
revision.search_space_id = 2
revision.workspace_id = 2
revision.content_before = content_before
revision.title_before = title_before
revision.folder_id_before = folder_id_before
@ -130,7 +130,7 @@ def _folder_revision(
revision = MagicMock()
revision.id = 200
revision.folder_id = folder_id
revision.search_space_id = 2
revision.workspace_id = 2
revision.name_before = name_before
revision.parent_id_before = parent_id_before
revision.position_before = position_before