Skip to content

Commit 95ad5f2

Browse files
ptiurinptiurin
committed
improve tests
1 parent 91ae327 commit 95ad5f2

File tree

1 file changed

+283
-0
lines changed

1 file changed

+283
-0
lines changed

‎tests/unit/service/test_engine.py‎

Lines changed: 283 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from typing import Callable, Union
2+
from unittest.mock import MagicMock, patch
23

34
from httpx import Request
45
from pytest import mark, raises
@@ -13,6 +14,43 @@
1314
from tests.unit.service.conftest import get_objects_from_db_callback
1415

1516

17+
def create_mock_engine_with_status_transitions(mock_engine: Engine, statuses: list):
18+
"""
19+
Helper function to create a callback that simulates engine status transitions.
20+
21+
Args:
22+
mock_engine: The base engine object to use for creating responses
23+
statuses: List of EngineStatus values to cycle through on subsequent calls
24+
25+
Returns:
26+
A callback function that can be used with HTTPXMock
27+
"""
28+
call_count = [0]
29+
30+
def get_engine_callback_with_transitions(request: Request) -> Response:
31+
# Return different statuses based on call count
32+
current_status = statuses[min(call_count[0], len(statuses) - 1)]
33+
call_count[0] += 1
34+
35+
engine_data = Engine(
36+
name=mock_engine.name,
37+
region=mock_engine.region,
38+
spec=mock_engine.spec,
39+
scale=mock_engine.scale,
40+
current_status=current_status,
41+
version=mock_engine.version,
42+
endpoint=mock_engine.endpoint,
43+
warmup=mock_engine.warmup,
44+
auto_stop=mock_engine.auto_stop,
45+
type=mock_engine.type,
46+
_database_name=mock_engine._database_name,
47+
_service=None,
48+
)
49+
return get_objects_from_db_callback([engine_data])(request)
50+
51+
return get_engine_callback_with_transitions
52+
53+
1654
def test_engine_create(
1755
httpx_mock: HTTPXMock,
1856
engine_name: str,
@@ -276,3 +314,248 @@ def test_engine_instantiation_with_different_configurations(
276314
assert engine.name == "test_engine"
277315
assert engine.region == "us-east-1"
278316
assert engine.scale == 2
317+
318+
319+
@patch("time.sleep")
320+
@patch("time.time")
321+
def test_engine_start_waits_for_draining_to_stop(
322+
mock_time: MagicMock,
323+
mock_sleep: MagicMock,
324+
httpx_mock: HTTPXMock,
325+
resource_manager: ResourceManager,
326+
mock_engine: Engine,
327+
system_engine_no_db_query_url: str,
328+
):
329+
"""
330+
Test that start() waits for an engine in DRAINING state to become STOPPED
331+
before proceeding with the start operation.
332+
"""
333+
# Set up time mock to avoid timeout - return incrementing values
334+
mock_time.return_value = 0 # Always return early time to avoid timeout
335+
336+
# Set up mock responses: DRAINING -> STOPPED -> STOPPED (after start command)
337+
callback = create_mock_engine_with_status_transitions(
338+
mock_engine,
339+
[
340+
EngineStatus.DRAINING, # Initial state
341+
EngineStatus.STOPPED, # After first refresh in _wait_for_start_stop
342+
EngineStatus.STOPPED, # After start command, final refresh
343+
],
344+
)
345+
346+
httpx_mock.add_callback(
347+
callback, url=system_engine_no_db_query_url, is_reusable=True
348+
)
349+
350+
# Set up the engine with proper service
351+
mock_engine._service = resource_manager.engines
352+
353+
# Call start method
354+
result = mock_engine.start()
355+
356+
# Verify that sleep was called (indicating it waited for DRAINING state)
357+
mock_sleep.assert_called_with(5)
358+
359+
# Verify the engine is returned
360+
assert result is mock_engine
361+
assert result.current_status == EngineStatus.STOPPED
362+
363+
364+
@patch("time.sleep")
365+
@patch("time.time")
366+
def test_engine_start_waits_for_stopping_to_stop(
367+
mock_time: MagicMock,
368+
mock_sleep: MagicMock,
369+
httpx_mock: HTTPXMock,
370+
resource_manager: ResourceManager,
371+
mock_engine: Engine,
372+
system_engine_no_db_query_url: str,
373+
):
374+
"""
375+
Test that start() waits for an engine in STOPPING state to become STOPPED
376+
before proceeding with the start operation.
377+
"""
378+
# Set up time mock to avoid timeout
379+
mock_time.return_value = 0 # Always return early time to avoid timeout
380+
381+
# Set up mock responses: STOPPING -> STOPPED -> STOPPED (after start command)
382+
callback = create_mock_engine_with_status_transitions(
383+
mock_engine,
384+
[
385+
EngineStatus.STOPPING, # Initial state
386+
EngineStatus.STOPPED, # After first refresh in _wait_for_start_stop
387+
EngineStatus.STOPPED, # After start command, final refresh
388+
],
389+
)
390+
391+
httpx_mock.add_callback(
392+
callback, url=system_engine_no_db_query_url, is_reusable=True
393+
)
394+
395+
# Set up the engine with proper service
396+
mock_engine._service = resource_manager.engines
397+
398+
# Call start method
399+
result = mock_engine.start()
400+
401+
# Verify that sleep was called (indicating it waited for STOPPING state)
402+
mock_sleep.assert_called_with(5)
403+
404+
# Verify the engine is returned
405+
assert result is mock_engine
406+
assert result.current_status == EngineStatus.STOPPED
407+
408+
409+
@patch("time.sleep")
410+
@patch("time.time")
411+
def test_engine_stop_waits_for_draining_to_stop(
412+
mock_time: MagicMock,
413+
mock_sleep: MagicMock,
414+
httpx_mock: HTTPXMock,
415+
resource_manager: ResourceManager,
416+
mock_engine: Engine,
417+
system_engine_no_db_query_url: str,
418+
):
419+
"""
420+
Test that stop() waits for an engine in DRAINING state to finish draining
421+
before proceeding with the stop operation.
422+
"""
423+
# Set up time mock to avoid timeout
424+
mock_time.return_value = 0 # Always return early time to avoid timeout
425+
426+
# Set up mock responses: DRAINING -> RUNNING -> STOPPED (after stop command)
427+
callback = create_mock_engine_with_status_transitions(
428+
mock_engine,
429+
[
430+
EngineStatus.DRAINING, # Initial state
431+
EngineStatus.RUNNING, # After first refresh in _wait_for_start_stop
432+
EngineStatus.STOPPED, # After stop command, final refresh
433+
],
434+
)
435+
436+
httpx_mock.add_callback(
437+
callback, url=system_engine_no_db_query_url, is_reusable=True
438+
)
439+
440+
# Set up the engine with proper service
441+
mock_engine._service = resource_manager.engines
442+
443+
# Call stop method
444+
result = mock_engine.stop()
445+
446+
# Verify that sleep was called (indicating it waited for DRAINING state)
447+
mock_sleep.assert_called_with(5)
448+
449+
# Verify the engine is returned
450+
assert result is mock_engine
451+
assert result.current_status == EngineStatus.STOPPED
452+
453+
454+
@patch("time.sleep")
455+
@patch("time.time")
456+
def test_engine_wait_for_start_stop_timeout(
457+
mock_time: MagicMock,
458+
mock_sleep: MagicMock,
459+
httpx_mock: HTTPXMock,
460+
resource_manager: ResourceManager,
461+
mock_engine: Engine,
462+
system_engine_no_db_query_url: str,
463+
):
464+
"""
465+
Test that _wait_for_start_stop raises TimeoutError when engine stays in
466+
transitional state too long.
467+
"""
468+
# Mock time.time to simulate timeout using a function that tracks calls
469+
call_count = [0]
470+
471+
def mock_time_function():
472+
call_count[0] += 1
473+
# Return normal time for first few calls, then timeout for _wait_for_start_stop
474+
if call_count[0] <= 5:
475+
return 0 # Early time
476+
else:
477+
return 3601 # Past timeout
478+
479+
mock_time.side_effect = mock_time_function
480+
481+
def get_engine_callback_always_starting(request: Request) -> Response:
482+
# Always return STARTING to simulate stuck state
483+
engine_data = Engine(
484+
name=mock_engine.name,
485+
region=mock_engine.region,
486+
spec=mock_engine.spec,
487+
scale=mock_engine.scale,
488+
current_status=EngineStatus.STARTING, # Always starting
489+
version=mock_engine.version,
490+
endpoint=mock_engine.endpoint,
491+
warmup=mock_engine.warmup,
492+
auto_stop=mock_engine.auto_stop,
493+
type=mock_engine.type,
494+
_database_name=mock_engine._database_name,
495+
_service=None,
496+
)
497+
return get_objects_from_db_callback([engine_data])(request)
498+
499+
httpx_mock.add_callback(
500+
get_engine_callback_always_starting,
501+
url=system_engine_no_db_query_url,
502+
is_reusable=True,
503+
)
504+
505+
# Set up the engine with proper service
506+
mock_engine._service = resource_manager.engines
507+
508+
# Call start method and expect TimeoutError
509+
with raises(TimeoutError, match="Excedeed timeout of 3600s waiting for.*starting"):
510+
mock_engine.start()
511+
512+
513+
@patch("time.sleep")
514+
@patch("time.time")
515+
def test_engine_start_already_running_no_wait(
516+
mock_time: MagicMock,
517+
mock_sleep: MagicMock,
518+
httpx_mock: HTTPXMock,
519+
resource_manager: ResourceManager,
520+
mock_engine: Engine,
521+
system_engine_no_db_query_url: str,
522+
):
523+
"""
524+
Test that start() doesn't wait when engine is already RUNNING.
525+
"""
526+
# Mock time to avoid any timeout issues
527+
mock_time.return_value = 0
528+
529+
def get_engine_callback_running(request: Request) -> Response:
530+
engine_data = Engine(
531+
name=mock_engine.name,
532+
region=mock_engine.region,
533+
spec=mock_engine.spec,
534+
scale=mock_engine.scale,
535+
current_status=EngineStatus.RUNNING,
536+
version=mock_engine.version,
537+
endpoint=mock_engine.endpoint,
538+
warmup=mock_engine.warmup,
539+
auto_stop=mock_engine.auto_stop,
540+
type=mock_engine.type,
541+
_database_name=mock_engine._database_name,
542+
_service=None,
543+
)
544+
return get_objects_from_db_callback([engine_data])(request)
545+
546+
httpx_mock.add_callback(
547+
get_engine_callback_running, url=system_engine_no_db_query_url, is_reusable=True
548+
)
549+
550+
# Set up the engine with proper service
551+
mock_engine._service = resource_manager.engines
552+
553+
# Call start method
554+
result = mock_engine.start()
555+
556+
# Verify that no sleep was called (no waiting happened)
557+
mock_sleep.assert_not_called()
558+
559+
# Verify the engine is returned
560+
assert result is mock_engine
561+
assert result.current_status == EngineStatus.RUNNING

0 commit comments

Comments
 (0)