Use Base test

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

Author: Yikun

tests/ut/base.py

@@ -1,3 +1,18 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 import unittest
 
 from vllm_ascend.utils import adapt_patch

tests/ut/distributed/test_parallel_state.py

@@ -1,16 +1,31 @@
-import unittest
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 from unittest.mock import MagicMock, patch
 
 import pytest
 from vllm.distributed.parallel_state import GroupCoordinator
 
 import vllm_ascend
+from tests.ut.base import TestBase
 from vllm_ascend.distributed.parallel_state import (
     destory_ascend_model_parallel, get_ep_group, get_etp_group,
     init_ascend_model_parallel, model_parallel_initialized)
 
 
-class TestParallelState(unittest.TestCase):
+class TestParallelState(TestBase):
 
     @patch('vllm_ascend.distributed.parallel_state._EP',
            new_callable=lambda: MagicMock(spec=GroupCoordinator))

tests/ut/ops/test_expert_load_balancer.py

@@ -1,14 +1,29 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 # fused moe ops test will hit the infer_schema error, we need add the patch
 # here to make the test pass.
 import vllm_ascend.patch.worker.patch_common.patch_utils  # type: ignore[import]  # isort: skip  # noqa
 
 import json
-import unittest
 from typing import List, TypedDict
 from unittest import mock
 
 import torch
 
+from tests.ut.base import TestBase
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
 
 
@@ -47,7 +62,7 @@ class MockData(TypedDict):
 }
 
 
-class TestExpertLoadBalancer(unittest.TestCase):
+class TestExpertLoadBalancer(TestBase):
 
     def setUp(self):
         json_file = "expert_map.json"

tests/ut/patch/worker/patch_common/test_patch_distributed.py

@@ -1,12 +1,27 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 from tests.ut.base import TestBase
+from vllm_ascend.patch.worker.patch_common.patch_distributed import \
+    GroupCoordinatorPatch
+
+# import GroupCoordinator after GroupCoordinatorPatch to make base work
+from vllm.distributed.parallel_state import GroupCoordinator  # noqa isort:skip
 
 
 class TestPatchDistributed(TestBase):
 
     def test_GroupCoordinator_patched(self):
-        from vllm.distributed.parallel_state import GroupCoordinator
-
-        from vllm_ascend.patch.worker.patch_common.patch_distributed import \
-            GroupCoordinatorPatch
-
         self.assertIs(GroupCoordinator, GroupCoordinatorPatch)

tests/ut/patch/worker/patch_common/test_patch_sampler.py

@@ -1,20 +1,32 @@
-import importlib
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 import os
-import unittest
 from unittest import mock
 
 import torch
 from vllm.v1.sample.ops import topk_topp_sampler
 
+from tests.ut.base import TestBase
+
 
-class TestTopKTopPSamplerOptimize(unittest.TestCase):
+class TestTopKTopPSamplerOptimize(TestBase):
 
     @mock.patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_TOPK_OPTIMIZE": "1"})
     @mock.patch("torch_npu.npu_top_k_top_p")
     def test_npu_topk_topp_called_when_optimized(self, mock_npu_op):
-        import vllm_ascend.patch.worker.patch_common.patch_sampler
-        importlib.reload(vllm_ascend.patch.worker.patch_common.patch_sampler)
-
         mock_npu_op.return_value = (torch.randn(1, 3))
         sampler = topk_topp_sampler.TopKTopPSampler()
 

tests/ut/test_ascend_config.py

@@ -1,16 +1,31 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 import os
-import unittest
 from unittest import mock
 
 from transformers import PretrainedConfig
 from vllm.config import ModelConfig, VllmConfig
 
+from tests.ut.base import TestBase
 from vllm_ascend.ascend_config import (check_ascend_config,
                                        clear_ascend_config, get_ascend_config,
                                        init_ascend_config)
 
 
-class TestAscendConfig(unittest.TestCase):
+class TestAscendConfig(TestBase):
 
     @staticmethod
     def _clean_up_ascend_config(func):

tests/ut/test_utils.py

@@ -1,3 +1,18 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+
 import math
 import os
 import unittest
@@ -8,10 +23,11 @@
 from vllm.config import (CompilationConfig, ModelConfig, ParallelConfig,
                          VllmConfig)
 
+from tests.ut.base import TestBase
 from vllm_ascend import utils
 
 
-class TestUtils(unittest.TestCase):
+class TestUtils(TestBase):
 
     def test_is_310p(self):
         utils._IS_310P = None