From 38a32186220ee57f9a6da51c6b05d0e62b61bc35 Mon Sep 17 00:00:00 2001 From: Yikun Jiang Date: Sun, 6 Jul 2025 10:49:30 +0800 Subject: [PATCH] [FOLLOWUP] Use base test Signed-off-by: Yikun Jiang --- tests/e2e/singlecard/ops/test_fused_moe.py | 4 ---- tests/ut/base.py | 8 ++------ .../ut/distributed/kv_transfer/test_simple_buffer.py | 4 ++-- .../distributed/kv_transfer/test_simple_connector.py | 4 ++-- tests/ut/distributed/kv_transfer/test_simple_pipe.py | 4 ++-- tests/ut/ops/test_rotary_embedding.py | 11 +++++------ tests/ut/test_utils.py | 3 +-- tests/ut/worker/test_input_batch.py | 5 ++--- tests/ut/worker/test_pooling_model_runner.py | 3 ++- 9 files changed, 18 insertions(+), 28 deletions(-) diff --git a/tests/e2e/singlecard/ops/test_fused_moe.py b/tests/e2e/singlecard/ops/test_fused_moe.py index 78c0d88684..76f1846382 100644 --- a/tests/e2e/singlecard/ops/test_fused_moe.py +++ b/tests/e2e/singlecard/ops/test_fused_moe.py @@ -19,10 +19,6 @@ Run `pytest tests/ops/test_fused_moe.py`. """ -# 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 pytest import torch from vllm.model_executor.layers.activation import SiluAndMul diff --git a/tests/ut/base.py b/tests/ut/base.py index e34f175935..fb1a625077 100644 --- a/tests/ut/base.py +++ b/tests/ut/base.py @@ -17,15 +17,11 @@ from vllm_ascend.utils import adapt_patch -# 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 - class TestBase(unittest.TestCase): - def setUp(self): + def __init__(self, *args, **kwargs): # adapt patch by default. adapt_patch(True) adapt_patch() - super().setUp() + super(TestBase, self).__init__(*args, **kwargs) diff --git a/tests/ut/distributed/kv_transfer/test_simple_buffer.py b/tests/ut/distributed/kv_transfer/test_simple_buffer.py index 6f90df923f..1ff81bc3b6 100644 --- a/tests/ut/distributed/kv_transfer/test_simple_buffer.py +++ b/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() diff --git a/tests/ut/distributed/kv_transfer/test_simple_connector.py b/tests/ut/distributed/kv_transfer/test_simple_connector.py index ac6c4d478d..2c81943a7c 100644 --- a/tests/ut/distributed/kv_transfer/test_simple_connector.py +++ b/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) diff --git a/tests/ut/distributed/kv_transfer/test_simple_pipe.py b/tests/ut/distributed/kv_transfer/test_simple_pipe.py index efd6eddea8..ccc984b62a 100644 --- a/tests/ut/distributed/kv_transfer/test_simple_pipe.py +++ b/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): diff --git a/tests/ut/ops/test_rotary_embedding.py b/tests/ut/ops/test_rotary_embedding.py index 91c2ad40df..3b388e09f2 100644 --- a/tests/ut/ops/test_rotary_embedding.py +++ b/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) diff --git a/tests/ut/test_utils.py b/tests/ut/test_utils.py index 5ddc59dea5..346e7c497f 100644 --- a/tests/ut/test_utils.py +++ b/tests/ut/test_utils.py @@ -15,7 +15,6 @@ import math import os -import unittest from threading import Lock from unittest import mock @@ -302,7 +301,7 @@ def test_torchair_cache_dir(self): "Delete kv cache bytes cache dir failed") -class TestProfileExecuteDuration(unittest.TestCase): +class TestProfileExecuteDuration(TestBase): def setUp(self): utils.ProfileExecuteDuration._instance = None diff --git a/tests/ut/worker/test_input_batch.py b/tests/ut/worker/test_input_batch.py index cbfd67f0a2..7baee71226 100644 --- a/tests/ut/worker/test_input_batch.py +++ b/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 diff --git a/tests/ut/worker/test_pooling_model_runner.py b/tests/ut/worker/test_pooling_model_runner.py index 28a0a7d3c6..7c5e6f59a0 100644 --- a/tests/ut/worker/test_pooling_model_runner.py +++ b/tests/ut/worker/test_pooling_model_runner.py @@ -7,11 +7,12 @@ from vllm.pooling_params import PoolingParams from vllm.sequence import SequenceData, SequenceGroupMetadata +from tests.ut.base import TestBase from vllm_ascend.worker.pooling_model_runner import ( ModelInputForNPUWithPoolingMetadata, NPUPoolingModelRunner) -class TestPoolingModelRunner(unittest.TestCase): +class TestPoolingModelRunner(TestBase): """Unit tests for the NPUPoolingModelRunner class.""" def _create_model_runner(self, model: str, *args,