
local lu = require 'luaunit'
local imu_cmd = require 'imu'
local ipmi = require 'ipmi'
local bs = require 'mc.bitstring'

TestNPUImuCmd = {}

function TestNPUImuCmd:test_pcie_info()
    local tuple_info_rsp = bs.new([[<<
        head:3/unit:8,
        total_length:4/little-unit:8,
        respond_length:4/little-unit:8,
        vendor_id:4/unit:8,
        device_id:4/unit:8,
        sub_vendor_id:4/unit:8,
        sub_device_id:4/unit:8
    >>]])
    local pcie_info_rsp = bs.new([[<<
        head:3/unit:8,
        total_length:4/little-unit:8,
        respond_length:4/little-unit:8,
        pcie_lanes:4/unit:8,
        pcie_rates:4/unit:8,
        dev_bus:4/unit:8,
        dev_device:4/unit:8,
        dev_function:4/unit:8
    >>]])
    local test_data = pcie_info_rsp:pack({
        head = 1,
        total_length = 1,
        respond_length = 1,
        pcie_lanes = 0,
        pcie_rates = 0, 
        dev_bus = 0,
        dev_device = 0,
        dev_function = 0
    })
    ipmi.request = function()
        return 0, test_data
    end
    local ok = pcall(function()
        imu_cmd.get_npu_pcie_info(0, 0)
    end)
    lu.assertEquals(ok, true)

    local test_data_tuple = tuple_info_rsp:pack({
        head = 1,
        total_length = 1,
        respond_length = 1,
        vendor_id = 0,
        device_id = 0, 
        sub_vendor_id = 0,
        sub_device_id = 0
    })
    ipmi.request = function()
        return 0, test_data_tuple
    end
    ok = pcall(function()
        imu_cmd.get_npu_pcie_sub_vid_did_info(0, 0)
    end)
    lu.assertEquals(ok, true)
end

function test_get_info_from_pmu()
    local c_bisstring = package.loaded['mc.bitstring']
    local mock_bs = {
        new = function()
            return {
                pack = function()
                end,
                unpack = function()
                    return {
                        info_byte = 0
                    }
                end
            }
        end
    }

    local pcie_info = {
        is_local = 1,
        cpu_id = 1,
        address = 1,
        function_num = 1,
        device_num = 1,
        bus_num = 1,
        read_length = 4
    }

    package.loaded['mc.bitstring'] = nil
    package.loaded['mc.bitstring'] = mock_bs
    package.loaded['imu'] = nil
    local mock_imu = require ('imu.init')
    local ok = mock_imu.get_info_from_pmu(1, {}, pcie_info)
    lu.assertEquals(ok, {0})

    local c_request = ipmi.request
    ipmi.request = function()
        return 0xff
    end
    ok = mock_imu.get_info_from_pmu(1, {}, pcie_info)
    lu.assertEquals(ok, nil)
    ipmi.request = c_request
    package.loaded['mc.bitstring'] = c_bisstring
end