Skip to content

Commit 4f36c2c

Browse files
jopemachineclaude
andcommitted
test(BA-6953): add component tests for the kernel scheduling-history v2 SDK client
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 9594aab commit 4f36c2c

3 files changed

Lines changed: 416 additions & 2 deletions

File tree

changes/12997.test.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add component tests for the kernel scheduling-history v2 SDK client.

tests/component/scheduling_history/conftest.py

Lines changed: 215 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,25 @@
11
from __future__ import annotations
22

3-
from unittest.mock import AsyncMock
3+
import uuid
4+
from collections.abc import AsyncIterator
5+
from dataclasses import dataclass
6+
from datetime import UTC, datetime, timedelta
7+
from decimal import Decimal
8+
from typing import TYPE_CHECKING
9+
from unittest.mock import AsyncMock, MagicMock
410

511
import pytest
12+
import sqlalchemy as sa
13+
import yarl
614

15+
from ai.backend.client.v2.auth import HMACAuth
16+
from ai.backend.client.v2.config import ClientConfig
17+
from ai.backend.client.v2.v2_registry import V2ClientRegistry
18+
from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName
19+
from ai.backend.common.types import ResourceSlot
720
from ai.backend.manager.actions.validators import ActionValidators
821
from ai.backend.manager.actions.validators.rbac import RBACValidators
22+
from ai.backend.manager.api.adapters.scheduling_history.adapter import SchedulingHistoryAdapter
923
from ai.backend.manager.api.rest.routing import RouteRegistry
1024

1125
# Statically imported so that Pants includes these modules in the test PEX.
@@ -16,12 +30,23 @@
1630
register_scheduling_history_routes,
1731
)
1832
from ai.backend.manager.api.rest.types import RouteDeps
33+
from ai.backend.manager.api.rest.v2.scheduling_history.handler import V2SchedulingHistoryHandler
34+
from ai.backend.manager.api.rest.v2.scheduling_history.registry import (
35+
register_v2_scheduling_history_routes,
36+
)
37+
from ai.backend.manager.models.kernel import KernelRow
38+
from ai.backend.manager.models.scheduling_history.row import KernelSchedulingHistoryRow
39+
from ai.backend.manager.models.session import SessionRow
1940
from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
2041
from ai.backend.manager.repositories.scheduling_history.repository import (
2142
SchedulingHistoryRepository,
2243
)
2344
from ai.backend.manager.services.scheduling_history.processors import SchedulingHistoryProcessors
2445
from ai.backend.manager.services.scheduling_history.service import SchedulingHistoryService
46+
from ai.backend.testutils.fixtures import DomainFixtureData
47+
48+
if TYPE_CHECKING:
49+
from tests.component.conftest import ServerInfo, UserFixtureData
2550

2651

2752
@pytest.fixture()
@@ -39,15 +64,203 @@ def scheduling_history_processors(
3964
)
4065

4166

67+
@pytest.fixture()
68+
def scheduling_history_adapter(
69+
scheduling_history_processors: SchedulingHistoryProcessors,
70+
) -> SchedulingHistoryAdapter:
71+
"""Build an adapter wired only with scheduling-history processors.
72+
73+
Every call site in the adapter goes through ``self._processors.scheduling_history``,
74+
so a MagicMock backing object with that attribute set is sufficient.
75+
"""
76+
processors = MagicMock()
77+
processors.scheduling_history = scheduling_history_processors
78+
return SchedulingHistoryAdapter(processors)
79+
80+
4281
@pytest.fixture()
4382
def server_module_registries(
4483
route_deps: RouteDeps,
4584
scheduling_history_processors: SchedulingHistoryProcessors,
85+
scheduling_history_adapter: SchedulingHistoryAdapter,
4686
) -> list[RouteRegistry]:
47-
"""Load only the modules required for scheduling-history domain tests."""
87+
"""Load both v1 and v2 scheduling-history route trees."""
88+
v2_registry = RouteRegistry.create("v2", route_deps.cors_options)
89+
v2_registry.add_subregistry(
90+
register_v2_scheduling_history_routes(
91+
V2SchedulingHistoryHandler(adapter=scheduling_history_adapter),
92+
route_deps,
93+
)
94+
)
4895
return [
4996
register_scheduling_history_routes(
5097
SchedulingHistoryHandler(scheduling_history=scheduling_history_processors),
5198
route_deps,
5299
),
100+
v2_registry,
101+
]
102+
103+
104+
@pytest.fixture()
105+
async def admin_v2_registry(
106+
server: ServerInfo,
107+
admin_user_fixture: UserFixtureData,
108+
) -> AsyncIterator[V2ClientRegistry]:
109+
registry = await V2ClientRegistry.create(
110+
ClientConfig(endpoint=yarl.URL(server.url)),
111+
HMACAuth(
112+
access_key=admin_user_fixture.keypair.access_key,
113+
secret_key=admin_user_fixture.keypair.secret_key,
114+
),
115+
)
116+
try:
117+
yield registry
118+
finally:
119+
await registry.close()
120+
121+
122+
@pytest.fixture()
123+
async def user_v2_registry(
124+
server: ServerInfo,
125+
regular_user_fixture: UserFixtureData,
126+
) -> AsyncIterator[V2ClientRegistry]:
127+
registry = await V2ClientRegistry.create(
128+
ClientConfig(endpoint=yarl.URL(server.url)),
129+
HMACAuth(
130+
access_key=regular_user_fixture.keypair.access_key,
131+
secret_key=regular_user_fixture.keypair.secret_key,
132+
),
133+
)
134+
try:
135+
yield registry
136+
finally:
137+
await registry.close()
138+
139+
140+
@dataclass(frozen=True)
141+
class KernelHistorySeed:
142+
"""Identifiers of the seeded kernels and the history rows attached to them."""
143+
144+
session_id: uuid.UUID
145+
kernel_id: uuid.UUID
146+
other_kernel_id: uuid.UUID
147+
history_ids: list[uuid.UUID]
148+
other_history_id: uuid.UUID
149+
150+
151+
@pytest.fixture()
152+
async def kernel_history_seed(
153+
database_engine: ExtendedAsyncSAEngine,
154+
domain_fixture: DomainFixtureData,
155+
group_fixture: uuid.UUID,
156+
scaling_group_name: ResourceGroupName,
157+
scaling_group_id: ResourceGroupID,
158+
admin_user_fixture: UserFixtureData,
159+
) -> AsyncIterator[KernelHistorySeed]:
160+
"""Seed one session with two kernels: three history rows on the first, one on the second.
161+
162+
``kernel_scheduling_history`` has no writer yet (BA-6852), so the rows go in directly.
163+
The scoped endpoint runs an existence check against ``kernels``, hence the real kernel rows.
164+
"""
165+
session_id = uuid.uuid4()
166+
kernel_id = uuid.uuid4()
167+
other_kernel_id = uuid.uuid4()
168+
slots = ResourceSlot({"cpu": Decimal("1"), "mem": Decimal("1073741824")})
169+
now = datetime.now(tz=UTC)
170+
171+
# created_at is staggered so ordering and cursor pagination are deterministic.
172+
phases = [
173+
("PREPARING", "PREPARING", "PULLING", "SUCCESS", 1),
174+
("PULLING", "PULLING", "PREPARED", "SUCCESS", 2),
175+
("RUNNING", "CREATING", "RUNNING", "FAILURE", 3),
53176
]
177+
history_ids = [uuid.uuid4() for _ in phases]
178+
other_history_id = uuid.uuid4()
179+
180+
async with database_engine.begin_session() as db_sess:
181+
db_sess.add(
182+
SessionRow(
183+
id=session_id,
184+
name=f"session-{session_id.hex[:8]}",
185+
domain_name=domain_fixture.domain_name,
186+
domain_id=domain_fixture.domain_id,
187+
group_id=group_fixture,
188+
user_uuid=admin_user_fixture.user_uuid,
189+
scaling_group_name=scaling_group_name,
190+
resource_group_id=scaling_group_id,
191+
occupying_slots=slots,
192+
requested_slots=slots,
193+
)
194+
)
195+
await db_sess.flush()
196+
db_sess.add_all([
197+
KernelRow(
198+
id=kid,
199+
session_id=session_id,
200+
domain_name=domain_fixture.domain_name,
201+
group_id=group_fixture,
202+
user_uuid=admin_user_fixture.user_uuid,
203+
occupied_slots=slots,
204+
requested_slots=slots,
205+
repl_in_port=0,
206+
repl_out_port=0,
207+
stdin_port=0,
208+
stdout_port=0,
209+
scaling_group=scaling_group_name,
210+
resource_group_id=scaling_group_id,
211+
)
212+
for kid in (kernel_id, other_kernel_id)
213+
])
214+
await db_sess.flush()
215+
216+
for offset, (hid, phase_row) in enumerate(zip(history_ids, phases, strict=False)):
217+
phase, from_status, to_status, result, attempts = phase_row
218+
db_sess.add(
219+
KernelSchedulingHistoryRow(
220+
id=hid,
221+
kernel_id=kernel_id,
222+
session_id=session_id,
223+
phase=phase,
224+
from_status=from_status,
225+
to_status=to_status,
226+
result=result,
227+
error_code=None if result == "SUCCESS" else "ERR_KERNEL_START",
228+
message=f"{phase} transition",
229+
attempts=attempts,
230+
created_at=now + timedelta(seconds=offset),
231+
updated_at=now + timedelta(seconds=offset),
232+
)
233+
)
234+
db_sess.add(
235+
KernelSchedulingHistoryRow(
236+
id=other_history_id,
237+
kernel_id=other_kernel_id,
238+
session_id=session_id,
239+
phase="PREPARING",
240+
from_status="PREPARING",
241+
to_status="PULLING",
242+
result="SUCCESS",
243+
error_code=None,
244+
message="other kernel transition",
245+
attempts=1,
246+
created_at=now,
247+
updated_at=now,
248+
)
249+
)
250+
251+
yield KernelHistorySeed(
252+
session_id=session_id,
253+
kernel_id=kernel_id,
254+
other_kernel_id=other_kernel_id,
255+
history_ids=history_ids,
256+
other_history_id=other_history_id,
257+
)
258+
259+
async with database_engine.begin() as conn:
260+
await conn.execute(
261+
sa.delete(KernelSchedulingHistoryRow).where(
262+
KernelSchedulingHistoryRow.session_id == session_id
263+
)
264+
)
265+
await conn.execute(sa.delete(KernelRow).where(KernelRow.session_id == session_id))
266+
await conn.execute(sa.delete(SessionRow).where(SessionRow.id == session_id))

0 commit comments

Comments
 (0)