-- Copyright (c) 2026 Huawei Technologies Co., Ltd.
-- openUBMC is licensed under Mulan PSL v2.
-- You can use this software according to the terms and conditions of the Mulan PSL v2.
-- You may obtain a copy of Mulan PSL v2 at:
--          http://license.coscl.org.cn/MulanPSL2
-- THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
-- EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
-- MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
-- See the Mulan PSL v2 for more details.


local lu = require 'luaunit'
local pmu_cmd = require 'imu'
local PCI_CAP_ID_EXP<const> = 0x10

TEST_imu = {}

function TEST_imu:test_get_pcie_link_info()
    local bus = nil
    local pcie_info = {
        address = 1
    }

    local original_get_capid_addr = pmu_cmd.pcie_get_capid_addr
    local original_get_info_from_pmu = pmu_cmd.get_info_from_pmu

    pmu_cmd.pcie_get_capid_addr = function() end
    local ret = pmu_cmd.get_pcie_link_info(bus, pcie_info)
    lu.assertEquals(ret, nil)

    -- addr = 0x01, addr_link_cap = 0x0D, addr_link_status = 0x11
    pmu_cmd.pcie_get_capid_addr = function()
        return 0x01
    end
    pmu_cmd.get_info_from_pmu = function() end
    ret = pmu_cmd.get_pcie_link_info(bus, pcie_info)
    lu.assertEquals(ret, nil)

    pmu_cmd.get_info_from_pmu = function(_, pcie_info)
        if pcie_info['address'] == 0x0D then
            return {0x8A, 0x00} -- link_width_ability = 'X8', link_speed_capability = 'N/A'
        end
    end
    ret = pmu_cmd.get_pcie_link_info(bus, pcie_info)
    lu.assertEquals(ret, nil)

    pmu_cmd.get_info_from_pmu = function(_, pcie_info)
        if pcie_info['address'] == 0x0D then
            return {0xAA, 0x00} -- link_width_ability  = 'N/A', link_speed_capability = 'N/A'
        end
        if pcie_info['address'] == 0x11 then
            return {nil, nil, 0x05, 0x01} -- link_width = 'X16', link_speed = '32.0GT/s'
        end
    end
    ret = pmu_cmd.get_pcie_link_info(bus, pcie_info)
    lu.assertNotNil(ret)
    lu.assertEquals(ret.link_width_ability, 'N/A')
    lu.assertEquals(ret.link_speed_capability, 'N/A')
    lu.assertEquals(ret.link_width, 'X16')
    lu.assertEquals(ret.link_speed, '32.0GT/s')
    -- 还原
    pmu_cmd.pcie_get_capid_addr = original_get_capid_addr
    pmu_cmd.get_info_from_pmu = original_get_info_from_pmu
end

function TEST_imu:test_pcie_get_capid_addr()
    local bus = nil
    local pcie_info = {
        address = 1
    }
    local original_get_info_from_pmu = pmu_cmd.get_info_from_pmu

    pmu_cmd.get_info_from_pmu = function() end
    local ret = pmu_cmd.pcie_get_capid_addr(bus, pcie_info, PCI_CAP_ID_EXP)
    lu.assertEquals(ret, nil)

    local call_count = 0
    pmu_cmd.get_info_from_pmu = function()
        call_count = call_count + 1
        if call_count == 2 then -- 第二次失败
            return nil
        end
        return {0x0, nil}
    end
    ret = pmu_cmd.pcie_get_capid_addr(bus, pcie_info, PCI_CAP_ID_EXP)
    lu.assertEquals(ret, nil)
    lu.assertEquals(pcie_info.address, 0x0)

    pmu_cmd.get_info_from_pmu = function()
        return {0x0, 0x0}
    end
    ret = pmu_cmd.pcie_get_capid_addr(bus, pcie_info, PCI_CAP_ID_EXP)
    lu.assertEquals(ret, nil)

    pmu_cmd.get_info_from_pmu = function()
        return {0x0, 0x1}
    end
    ret = pmu_cmd.pcie_get_capid_addr(bus, pcie_info, PCI_CAP_ID_EXP)
    lu.assertEquals(ret, nil)
    lu.assertEquals(pcie_info.address, 0x1)

    pmu_cmd.get_info_from_pmu = function()
        return {PCI_CAP_ID_EXP, nil}
    end
    ret = pmu_cmd.pcie_get_capid_addr(bus, pcie_info, PCI_CAP_ID_EXP)
    lu.assertEquals(ret, 0x10)
    -- 还原
    pmu_cmd.get_info_from_pmu = original_get_info_from_pmu
end