56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.scripted_runtime import scheduler_hook
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
|
|
|
|
|
def _yielding_gen():
|
|
yield
|
|
yield
|
|
|
|
|
|
def _empty_gen():
|
|
return
|
|
yield # pragma: no cover — makes this a generator function
|
|
|
|
|
|
def _raising_gen():
|
|
raise ValueError("scripted-boom")
|
|
yield # pragma: no cover — makes this a generator function
|
|
|
|
|
|
class TestAdvanceGenerator(CustomTestCase):
|
|
def test_not_done_when_generator_yields(self):
|
|
done, exc_tb = scheduler_hook._advance_generator(_yielding_gen())
|
|
|
|
self.assertEqual((done, exc_tb), (False, None))
|
|
|
|
def test_done_without_traceback_on_stop_iteration(self):
|
|
done, exc_tb = scheduler_hook._advance_generator(_empty_gen())
|
|
|
|
self.assertEqual((done, exc_tb), (True, None))
|
|
|
|
def test_done_with_traceback_on_exception(self):
|
|
original = scheduler_hook.logger.exception
|
|
scheduler_hook.logger.exception = MagicMock()
|
|
try:
|
|
done, exc_tb = scheduler_hook._advance_generator(_raising_gen())
|
|
scheduler_hook.logger.exception.assert_called_once()
|
|
finally:
|
|
scheduler_hook.logger.exception = original
|
|
|
|
self.assertTrue(done)
|
|
self.assertIsNotNone(exc_tb)
|
|
self.assertIn("ValueError", exc_tb)
|
|
self.assertIn("scripted-boom", exc_tb)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|