# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.
from unittest.mock import MagicMock
import pytest

from mindspeed_rl.trainer.utils.parallel_state import get_pipeline_model_parallel_rank

from tests.test_tools.dist_test import DistributedTest


class TestParallelState(DistributedTest):
    world_size = 1
    is_dist_test = False

    def test_get_pipeline_model_parallel_rank(self):
        mpu = MagicMock()
        mpu.get_pipeline_model_parallel_rank.return_value = 1

        result = get_pipeline_model_parallel_rank(mpu, use_vllm=False)
        assert result == 1, "Parallel state method failed!"
