[FOLLOWUP] Use base test

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

Author: Yikun

tests/ut/distributed/kv_transfer/test_simple_buffer.py

@@ -1,9 +1,9 @@
-import unittest
 import zlib
 from unittest.mock import MagicMock
 
 import torch
 
+from tests.ut.base import TestBase
 from vllm_ascend.distributed.kv_transfer.simple_buffer import (SimpleBuffer,
                                                                int32_hash)
 
@@ -17,7 +17,7 @@ def __init__(self):
         self.deallocate_buffer = MagicMock()
 
 
-class TestSimpleBuffer(unittest.TestCase):
+class TestSimpleBuffer(TestBase):
 
     def setUp(self):
         self.pipe = MockSimplePipe()

tests/ut/distributed/kv_transfer/test_simple_connector.py

@@ -1,17 +1,17 @@
-import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
 from vllm.config import VllmConfig
 from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
 
+from tests.ut.base import TestBase
 from vllm_ascend.distributed.kv_transfer.simple_buffer import SimpleBuffer
 from vllm_ascend.distributed.kv_transfer.simple_connector import \
     SimpleConnector
 from vllm_ascend.distributed.kv_transfer.simple_pipe import SimplePipe
 
 
-class TestSimpleConnector(unittest.TestCase):
+class TestSimpleConnector(TestBase):
 
     def setUp(self):
         self.mock_pipe = MagicMock(spec=SimplePipe)

tests/ut/distributed/kv_transfer/test_simple_pipe.py

@@ -1,12 +1,12 @@
-import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
 
+from tests.ut.base import TestBase
 from vllm_ascend.distributed.kv_transfer.simple_pipe import SimplePipe
 
 
-class TestSimplePipe(unittest.TestCase):
+class TestSimplePipe(TestBase):
 
     @classmethod
     def _create_mock_config(self):

tests/ut/ops/test_rotary_embedding.py

@@ -1,5 +1,4 @@
 import math
-import unittest
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -12,7 +11,7 @@
                                               yarn_get_mscale)
 
 
-class TestCustomRotaryEmbeddingEnabled(unittest.TestCase):
+class TestCustomRotaryEmbeddingEnabled(TestBase):
 
     def setUp(self):
         # Common setup for tests
@@ -67,7 +66,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 +261,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 +271,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 +292,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

@@ -15,7 +15,6 @@
 
 import math
 import os
-import unittest
 from threading import Lock
 from unittest import mock
 
@@ -281,7 +280,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

@@ -1,11 +1,10 @@
-import unittest
-
 import numpy as np
 import torch
 from vllm.sampling_params import SamplingParams
 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 +23,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