[FOLLOWUP] Use base test

Signed-off-by: Yikun Jiang <yikunkero@gmail.com>

Author: Yikun

tests/ut/ops/test_rotary_embedding.py

@@ -12,7 +12,7 @@
                                               yarn_get_mscale)
 
 
-class TestCustomRotaryEmbeddingEnabled(unittest.TestCase):
+class TestCustomRotaryEmbeddingEnabled(TestBase):
 
     def setUp(self):
         # Common setup for tests
@@ -67,7 +67,7 @@ def test_custom_rotary_embedding_enabled(self):
             self.assertFalse(result)
 
 
-class TestRopeForwardOot(unittest.TestCase):
+class TestRopeForwardOot(TestBase):
 
     def setUp(self):
         # Common setup for tests
@@ -262,7 +262,7 @@ def test_native_rope_deepseek_forward_non_neox_style(
         assert k_pe.shape == key.shape
 
 
-class TestRotateHalf(unittest.TestCase):
+class TestRotateHalf(TestBase):
 
     def test_rotate_half_even_dim(self):
         # Test with even dimension
@@ -272,7 +272,7 @@ def test_rotate_half_even_dim(self):
         self.assertTrue(torch.allclose(result, expected))
 
 
-class TestYarnFindCorrectionDim(unittest.TestCase):
+class TestYarnFindCorrectionDim(TestBase):
 
     def test_basic_case(self):
         # Test with standard values
@@ -293,7 +293,7 @@ def test_basic_case(self):
         self.assertTrue(torch.allclose(result, expected))
 
 
-class TestYarnGetMscale(unittest.TestCase):
+class TestYarnGetMscale(TestBase):
 
     def test_scale_less_than_or_equal_1(self):
         self.assertEqual(yarn_get_mscale(scale=0.5), 1.0)

tests/ut/test_utils.py

@@ -19,6 +19,7 @@
 from threading import Lock
 from unittest import mock
 
+from tests.ut.base import TestBase
 import torch
 from vllm.config import (CompilationConfig, ModelConfig, ParallelConfig,
                          VllmConfig)
@@ -281,7 +282,7 @@ def test_update_aclgraph_sizes(self):
             len(test_vllm_config.compilation_config.cudagraph_capture_sizes))
 
 
-class TestProfileExecuteDuration(unittest.TestCase):
+class TestProfileExecuteDuration(TestBase):
 
     def setUp(self):
         utils.ProfileExecuteDuration._instance = None

tests/ut/worker/test_input_batch.py

@@ -6,6 +6,7 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.worker.block_table import MultiGroupBlockTable
 
+from tests.ut.base import TestBase
 from vllm_ascend.worker.npu_input_batch import CachedRequestState, InputBatch
 
 
@@ -24,7 +25,7 @@ def mock_cached_request_state(req_id="1", prompt=[1, 2, 3], output=[4, 5, 6]):
     )
 
 
-class TestInputBatch(unittest.TestCase):
+class TestInputBatch(TestBase):
 
     def setUp(self):
         self.max_num_reqs = 10