-- Copyright (c) 2024 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 sml = require 'sml'

local controller = {
    Id = 0,
    DevFunction = 1,
    SocketId = 1,
    DevDevice = 1,
    TypeId = 99,
    Segment = 1,
    SystemId = 1,
    DevBus = 1,
    Name = 'HI1880_SP686C_M_16i_2G',
    OOBSupport = 1,
    path = 'Controller_1_0001010102',
    CtrlOption1 = 2275649443,
    CtrlOption2 = 32647,
    CtrlOption3 = 0,
}

local function pd_equal(pd1, pd2)
    return pd1.controller_id == pd2.controller_id and pd1.device_id == pd2.device_id and
        pd1.slot_num == pd2.slot_num and pd1.enclosure_id == pd2.enclosure_id
end

function TestStorage:reset_pd_identify_func()
    ---@diagnostic disable-next-line: duplicate-set-field
    self.pd_identify_service.identify_pd_by_persistence_data = function(...)
        return
    end

    ---@diagnostic disable-next-line: duplicate-set-field
    self.pd_identify_service.update_identified_data = function(...)
        return
    end
end

function TestStorage:test_pd_identify()
    local drive_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0
    }
    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive_info)
    self.bus_monitor_service.on_os_state_changed:emit(true)

    self:reset_pd_identify_func()

    -- 这时候 pd_identify_service 的几个列表全部都是空的
    lu.assertEquals(#self.pd_identify_service.pd_list, 0)
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.identified_list, 0)

    -- 运行任务，这会触发 mctp_service::prepare 成功，跟着 object_manager 发送 on_add_object 信号，
    -- 最终 pd_identify_service 服务将收到添加 drive 对象的信号，同时生成配对任务
    self.tasks:run_until(function()
        return #self.pd_identify_service.drive_list == 1
    end)

    lu.assertNotIsNil(self.pd_identify_service.task)

    local pd1 = { controller_id = controller.Id, device_id = 1, slot_num = 2, enclosure_id = 3 }
    local pd2 = { controller_id = controller.Id, device_id = 2, slot_num = 2, enclosure_id = 3 }
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return { pd1, pd2 }
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end

    -- 运行任务，这会触发 controller_object 更新 pd_list，
    -- 最终 pd_identify_service 服务将收到添加 pd 对象的信号
    self.tasks:run_until(function()
        return #self.pd_identify_service.pd_list == 2
    end)
    lu.assertNotIsNil(self.pd_identify_service.task)

    -- 运行任务，这会触发 pd_identify_service.task 运行尝试定位，
    -- 但这次定位会失败，因为 drive:is_being_located 函数返回 false
    self.tasks:run_all_task()

    -- 再次运行任务，我们模拟 drive 点亮，使得 drive:is_being_located 返回 true，
    -- 这将触发一对 drive 和 pd 定位成功
    drive_info.LocateLed = 1
    drive_info.FaultLed = 0
    self.tasks:run_all_task()
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.pd_list, 1)         -- 前面模拟了 2 个 pd，现在定位成功了一个
    lu.assertEquals(#self.pd_identify_service.identified_list, 1) -- 已配对的{pd, drive}
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end

-- 测试硬盘点灯失败
function TestStorage:test_pd_link_abnormal()
    local drive_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0
    }
    
    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive_info)
    self.bus_monitor_service.on_os_state_changed:emit(true)
    -- 这时候 pd_identify_service 的几个列表全部都是空的
    lu.assertEquals(#self.pd_identify_service.pd_list, 0)
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.identified_list, 0)

    -- 运行任务，这会触发 mctp_service::prepare 成功，跟着 object_manager 发送 on_add_object 信号，
    -- 最终 pd_identify_service 服务将收到添加 drive 对象的信号
    self.tasks:run_until(function()
        return #self.pd_identify_service.drive_list == 1
    end)

    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return {}
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end

    -- 只要drive对象添加成功，就会运行pd_identify_service任务
    lu.assertNotIsNil(self.pd_identify_service.task)

    self.tasks:run_until(function()
        return self.pd_identify_service.add_pd_complete == true
    end)

    -- 但是没有pd对象，不点灯
    for i = 1, 60 do
        self.tasks:run_all_task()
    end

    lu.assertEquals(drive_info.Missing, 0)
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end

-- 测试 pd 被删除，模拟热插拔
function TestStorage:test_pd_identify_remove_pd()
    local drive_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0
    }

    self:reset_pd_identify_func()

    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive_info)
    local pd_list = {
        { controller_id = controller.Id, device_id = 1, slot_num = 2, enclosure_id = 3 },
        { controller_id = controller.Id, device_id = 2, slot_num = 2, enclosure_id = 3 }
    }
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return pd_list
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end
    self.bus_monitor_service.on_os_state_changed:emit(true)
    -- 运行任务直到配对成功
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 1
    end)
    local pd = self.pd_identify_service.identified_list[1].pd
    local drive = self.pd_identify_service.identified_list[1].drive
    lu.assertIsTrue(pd_equal(pd, pd_list[1]))
    lu.assertEquals(drive.Name, 'Disk1')

    -- 移除已配对的 pd
    table.remove(pd_list, 1)
    -- 运行任务，这会触发 controller 发送 pd 删除信号
    -- 接着 pd_identify_service 会解除配对关系，将 drive 重新放回未配对列表
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 0
    end)
    lu.assertEquals(#self.pd_identify_service.drive_list, 1) -- 重新将 drive 放回未配对列表
    lu.assertEquals(#self.pd_identify_service.pd_list, 1)

    -- 再次运行任务，由于没有将 drive.LocateLed 和 drive.FaultLed 复位，所以会与剩下的 pd 配对成功
    self.tasks:run_all_task()
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.pd_list, 0)
    lu.assertEquals(#self.pd_identify_service.identified_list, 1)
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end

-- 测试 drive 在位状态变化
function TestStorage:test_pd_identify_drive_not_presence()
    local drive_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0
    }

    self:reset_pd_identify_func()

    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive_info)
    self.bus_monitor_service.on_os_state_changed:emit(true)
    local pd_list = {
        { controller_id = controller.Id, device_id = 1, slot_num = 2, enclosure_id = 3 },
        { controller_id = controller.Id, device_id = 2, slot_num = 2, enclosure_id = 3 }
    }
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return pd_list
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end

    -- 运行任务直到配对成功
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 1
    end)
    local pd = self.pd_identify_service.identified_list[1].pd
    local drive = self.pd_identify_service.identified_list[1].drive
    lu.assertIsTrue(pd_equal(pd, pd_list[1]))
    lu.assertEquals(drive.Name, 'Disk1')

    -- 模拟 drive 不在位
    drive_info.Presence = 0
    -- 运行任务，这会触发 drive 发送不在位信号
    -- 接着 pd_identify_service 会解除配对关系
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 0
    end)
    self.tasks:run_all_task() -- 再次运行任务，更新pd列表
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.pd_list, 2)

    -- 模拟 drive 在位，drive 被放入待配对列表
    drive_info.Presence = 1
    self.tasks:run_until(function()
        return #self.pd_identify_service.drive_list == 1
    end)
    self.tasks:run_all_task() -- 再次运行任务，配对成功
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)
    lu.assertEquals(#self.pd_identify_service.pd_list, 1)
    lu.assertEquals(#self.pd_identify_service.identified_list, 1)
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end

-- 测试 pd 灯点亮后，循环检测 drive:is_being_located() 时发生 pd 热插拔
-- 这个用例的目的是验证定位算法在运行过程中协程被挂起时，发生 pd 插拔定位算法仍然能够正确运行
function TestStorage:test_pd_identify_remove_pd_1()
    local drive_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0,
        PassThrough = 0
    }

    self:reset_pd_identify_func()

    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive_info)
    local pd_list = {
        { controller_id = controller.Id, device_id = 1, slot_num = 2, enclosure_id = 3 },
        { controller_id = controller.Id, device_id = 2, slot_num = 2, enclosure_id = 3 }
    }
    self.bus_monitor_service.on_os_state_changed:emit(true)
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return pd_list
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end

    -- 在 map_allowed 函数的 tasks.sleep_ms 函数下钩子，模拟 sleep 时发生 pd 动态插拔
    local old_sleep_ms = self.tasks.sleep_ms
    local hook_sleep_ms = function(...)
        -- 先卸载 sleep_ms 钩子
        self.tasks.sleep_ms = old_sleep_ms

        -- 模拟移除一个 pd
        table.remove(pd_list, 1)
        self.tasks:run_all_task()

        old_sleep_ms(...)
    end
    local old_map_allowed = self.pd_identify_service.map_allowed
    self.pd_identify_service.map_allowed = function(...)
        self.tasks.sleep_ms = hook_sleep_ms                    -- 安装 sleep_ms 钩子
        self.pd_identify_service.map_allowed = old_map_allowed -- 卸载 map_allowed 钩子
        return old_map_allowed(...)
    end
    self.tasks:run_until(function()
        -- 运行直到 pd_list 的第一个 pd（device_id==1）被删除，只剩下 device_id==2 的 pd
        local pds = self.pd_identify_service.pd_list
        return #pds == 1 and pds[1].device_id == 2
    end)

    -- 插拔已经发生，本轮定位会失败，但不影响定位任务下一轮运行
    lu.assertEquals(#self.pd_identify_service.identified_list, 0)

    -- 卸载钩子后再次定位，这次能定位成功
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 1
    end)
    local identified = self.pd_identify_service.identified_list[1]
    lu.assertIsTrue(pd_equal(identified.pd, pd_list[1]))
    lu.assertEquals(identified.drive.Name, 'Disk1')
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end

local function init_data()
    local drive1_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 1,
        ObjectName = 'Drive1',
        LocateLed = 1,
        FaultLed = 0
    }
    local drive2_info = {
        HddBackplaneStartSlot = 0,
        Presence = 1,
        RelativeSlot = 2,
        ObjectName = 'Drive2',
        LocateLed = 1,
        FaultLed = 0
    }

    return drive1_info, drive2_info
end

-- 测试 pd 灯点亮后，循环检测 drive:is_being_located() 时发生 drive 热插拔
-- 这个用例的目的是验证定位算法在运行过程中协程被挂起时，发生 drive 插拔定位算法仍然能够正确运行
function TestStorage:test_pd_identify_remove_drive_1()
    local drive1_info, drive2_info = init_data()

    self:reset_pd_identify_func()
    self.object_manager.mc:add_object('Controller', controller, 1)
    self.object_manager.mc:add_object('Drive', drive1_info)
    self.object_manager.mc:add_object('Drive', drive2_info)
    local pd_list = {
        { controller_id = controller.Id, device_id = 1, slot_num = 2, enclosure_id = 3 },
        { controller_id = controller.Id, device_id = 2, slot_num = 2, enclosure_id = 3 }
    }
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_pd_list = function(...)
        return pd_list
    end
    ---@diagnostic disable-next-line: duplicate-set-field
    sml.get_ctrl_init_state = function(...)
        return 2
    end

    -- 在 map_allowed 函数的 tasks.sleep_ms 函数下钩子，模拟 sleep 时发生 pd 动态插拔
    local current_drive, current_pd
    local old_sleep_ms = self.tasks.sleep_ms
    local hook_sleep_ms = function(...)
        -- 模拟 drive 不在位
        current_drive.Presence = 0
        self.tasks:run_all_task()

        old_sleep_ms(...)
    end
    local old_map_allowed = self.pd_identify_service.map_allowed
    self.pd_identify_service.map_allowed = function(s, pd, drive)
        self.tasks.sleep_ms = hook_sleep_ms -- 安装 sleep_ms 钩子
        current_drive = drive
        current_pd = pd
        return old_map_allowed(s, pd, drive)
    end
    self.bus_monitor_service.on_os_state_changed:emit(true)
    self.tasks:run_until(function()
        -- 运行直到 drive1 不在位
        return drive1_info.Presence == 0 and drive2_info.Presence == 0
    end)

    -- drive 全部不在位，本轮定位会失败，self.pd_identify_service.drive_list 列表为空
    lu.assertEquals(#self.pd_identify_service.identified_list, 0)
    lu.assertEquals(#self.pd_identify_service.drive_list, 0)

    -- 卸载钩子后再次定位，设置 Disk2 在位，这次能定位成功
    drive1_info.Presence = 1
    hook_sleep_ms = function(...)
        if (current_pd.device_id == 1 and current_drive.Name == 'Disk1') or
            (current_pd.device_id == 2 and current_drive.Name == 'Disk2') then
            current_drive.LocateLed = 1
        else
            current_drive.LocateLed = 0
        end

        old_sleep_ms(...)
    end
    self.tasks.sleep_ms = old_sleep_ms
    self.tasks:run_until(function()
        return #self.pd_identify_service.identified_list == 1
    end)
    local identified = self.pd_identify_service.identified_list[1]
    lu.assertIsTrue(pd_equal(identified.pd, pd_list[1]))
    lu.assertEquals(identified.drive.Name, 'Disk1')
    -- 删除受模块管理的对象
    self.object_manager.mc:del_object('Controller', controller, 1)
    self.tasks:run_all_task()
end
