Add retry for flaky tests in CI (#4755)

test/lang/test_anthropic_backend.py

@@ -3,9 +3,10 @@ import unittest
 
 from sglang import Anthropic, set_default_backend
 from sglang.test.test_programs import test_mt_bench, test_stream
+from sglang.test.test_utils import CustomTestCase
 
 
-class TestAnthropicBackend(unittest.TestCase):
+class TestAnthropicBackend(CustomTestCase):
     backend = None
 
     @classmethod

test/lang/test_bind_cache.py

@@ -1,10 +1,10 @@
 import unittest
 
 import sglang as sgl
-from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST
+from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
 
 
-class TestBind(unittest.TestCase):
+class TestBind(CustomTestCase):
     backend = None
 
     @classmethod

test/lang/test_choices.py

@@ -7,6 +7,7 @@ from sglang.lang.choices import (
     token_length_normalized,
     unconditional_likelihood_normalized,
 )
+from sglang.test.test_utils import CustomTestCase
 
 MOCK_CHOICES_INPUT_DATA = {
     "choices": [
@@ -51,7 +52,7 @@ MOCK_CHOICES_INPUT_DATA = {
 }
 
 
-class TestChoices(unittest.TestCase):
+class TestChoices(CustomTestCase):
 
     def test_token_length_normalized(self):
         """Confirm 'antidisestablishmentarianism' is selected due to high confidences for

test/lang/test_litellm_backend.py

@@ -3,9 +3,10 @@ import unittest
 
 from sglang import LiteLLM, set_default_backend
 from sglang.test.test_programs import test_mt_bench, test_stream
+from sglang.test.test_utils import CustomTestCase
 
 
-class TestAnthropicBackend(unittest.TestCase):
+class TestAnthropicBackend(CustomTestCase):
     chat_backend = None
 
     @classmethod

test/lang/test_openai_backend.py

@@ -17,9 +17,10 @@ from sglang.test.test_programs import (
     test_stream,
     test_tool_use,
 )
+from sglang.test.test_utils import CustomTestCase
 
 
-class TestOpenAIBackend(unittest.TestCase):
+class TestOpenAIBackend(CustomTestCase):
     instruct_backend = None
     chat_backend = None
     chat_vision_backend = None

test/lang/test_srt_backend.py

@@ -22,10 +22,10 @@ from sglang.test.test_programs import (
     test_stream,
     test_tool_use,
 )
-from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST
+from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
 
 
-class TestSRTBackend(unittest.TestCase):
+class TestSRTBackend(CustomTestCase):
     backend = None
 
     @classmethod

test/lang/test_tracing.py

@@ -3,9 +3,10 @@ import unittest
 import sglang as sgl
 from sglang.lang.backend.base_backend import BaseBackend
 from sglang.lang.chat_template import get_chat_template
+from sglang.test.test_utils import CustomTestCase
 
 
-class TestTracing(unittest.TestCase):
+class TestTracing(CustomTestCase):
     def test_few_shot_qa(self):
         @sgl.function
         def few_shot_qa(s, question):

test/lang/test_vertexai_backend.py

@@ -10,9 +10,10 @@ from sglang.test.test_programs import (
     test_parallel_encoding,
     test_stream,
 )
+from sglang.test.test_utils import CustomTestCase
 
 
-class TestVertexAIBackend(unittest.TestCase):
+class TestVertexAIBackend(CustomTestCase):
     backend = None
 
     @classmethod