11from __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
511import 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
720from ai .backend .manager .actions .validators import ActionValidators
821from ai .backend .manager .actions .validators .rbac import RBACValidators
22+ from ai .backend .manager .api .adapters .scheduling_history .adapter import SchedulingHistoryAdapter
923from ai .backend .manager .api .rest .routing import RouteRegistry
1024
1125# Statically imported so that Pants includes these modules in the test PEX.
1630 register_scheduling_history_routes ,
1731)
1832from 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
1940from ai .backend .manager .models .utils import ExtendedAsyncSAEngine
2041from ai .backend .manager .repositories .scheduling_history .repository import (
2142 SchedulingHistoryRepository ,
2243)
2344from ai .backend .manager .services .scheduling_history .processors import SchedulingHistoryProcessors
2445from 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 ()
4382def 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