-- Copyright (c) 2026 Huawei Technologies Co., Ltd.
-- openUBMC is licensed under Mulan PSL v2.

local lu = require('luaunit')

local parallel_scheduler = require('upgrade.parallel_scheduler')
local parser_cfg = require('upgrade.parser_cfg')
TestParallelScheduler = {}

function TestParallelScheduler:test_new_and_capacity()
    local sched = parallel_scheduler.new(4)
    lu.assertEquals(sched:get_capacity(), 4)
    lu.assertEquals(sched:get_running(), 0)
    lu.assertEquals(sched:get_waiting(), 0)
end

function TestParallelScheduler:test_set_capacity()
    local sched = parallel_scheduler.new(2)
    sched:set_capacity(8)
    lu.assertEquals(sched:get_capacity(), 8)
end

function TestParallelScheduler:test_get_running_and_waiting()
    local sched = parallel_scheduler.new(2)
    lu.assertEquals(sched:get_running(), 0)
    lu.assertEquals(sched:get_waiting(), 0)
end

TestParserCfg = {}

function TestParserCfg:test_parse_dir()
    local dir, name = parser_cfg.parse_dir('/tmp/upgrade/Firmware0.hpm')
    lu.assertStrContains(dir, '/tmp/upgrade/')
    lu.assertEquals(name, 'Firmware0.hpm')
end

function TestParserCfg:test_get_cfgs_nil_path()
    local cfgs = parser_cfg.get_cfgs(nil)
    lu.assertEquals(next(cfgs), nil)
end

function TestParserCfg:test_get_cfgs_no_basic_version()
    local ini_parser = require('mc.v2_persistence')
    local orig_load = ini_parser.load_file
    ini_parser.load_file = function()
        return { Basic = {} }
    end
    local cfgs = parser_cfg.get_cfgs('/fake/path')
    lu.assertEquals(next(cfgs), nil)
    ini_parser.load_file = orig_load
end

function TestParserCfg:test_get_cfgs_valid()
    local ini_parser = require('mc.v2_persistence')
    local orig_load = ini_parser.load_file
    ini_parser.load_file = function()
        return {
            Basic = { Version = '1.0' },
            Firmware = { FileNum = '3' },
            Firmware1 = { ComponentID = '0x22', ComponentIDEx = '0xffffffff', Model = 'ModelA,ModelB' },
            Firmware2 = { ComponentID = '0x11', ComponentIDEx = '0x01', Model = 'ModelC' },
        }
    end
    local cfgs = parser_cfg.get_cfgs('/fake/path')
    lu.assertNotNil(cfgs.Firmware1)
    lu.assertEquals(cfgs.Firmware1.component_id, 0x22)
    lu.assertEquals(cfgs.Firmware1.component_idex, 0xffffffff)
    lu.assertEquals(cfgs.Firmware1.model[1], 'ModelA')
    lu.assertEquals(cfgs.Firmware1.model[2], 'ModelB')
    lu.assertNotNil(cfgs.Firmware2)
    ini_parser.load_file = orig_load
end

function TestParserCfg:test_get_cfgs_zero_file_num()
    local ini_parser = require('mc.v2_persistence')
    local orig_load = ini_parser.load_file
    ini_parser.load_file = function()
        return {
            Basic = { Version = '1.0' },
            Firmware = { FileNum = '0' },
        }
    end
    local cfgs = parser_cfg.get_cfgs('/fake/path')
    lu.assertEquals(next(cfgs), nil)
    ini_parser.load_file = orig_load
end

TestControllerUpgrade = {}

function TestControllerUpgrade:setUp()
    self.ctrl_upgrade = require('upgrade.controller_upgrade')
    self._orig_is_upgrading = self.ctrl_upgrade.is_controller_upgrading
    self._orig_cfgs = self.ctrl_upgrade.cfgs
    self._orig_upgrade_list = self.ctrl_upgrade.upgrade_list
    self._orig_before_version = self.ctrl_upgrade.before_upgrade_version
    self._orig_count = self.ctrl_upgrade.upgrade_controller_count
end

function TestControllerUpgrade:tearDown()
    if self._orig_is_upgrading ~= nil then
        self.ctrl_upgrade.is_controller_upgrading = self._orig_is_upgrading
    end
    if self._orig_cfgs ~= nil then
        self.ctrl_upgrade.cfgs = self._orig_cfgs
    end
    if self._orig_upgrade_list ~= nil then
        self.ctrl_upgrade.upgrade_list = self._orig_upgrade_list
    end
    if self._orig_before_version ~= nil then
        self.ctrl_upgrade.before_upgrade_version = self._orig_before_version
    end
    if self._orig_count ~= nil then
        self.ctrl_upgrade.upgrade_controller_count = self._orig_count
    end
end

-- 覆盖 init: 初始化成员变量
function TestControllerUpgrade:test_init()
    self.ctrl_upgrade.is_controller_upgrading = true
    self.ctrl_upgrade.cfgs = { fake = 1 }
    self.ctrl_upgrade:init()
    lu.assertFalse(self.ctrl_upgrade.is_controller_upgrading)
    lu.assertEquals(next(self.ctrl_upgrade.cfgs), nil)
    lu.assertEquals(self.ctrl_upgrade.before_upgrade_version, '')
    lu.assertEquals(self.ctrl_upgrade.upgrade_controller_count, 0)
end

-- 覆盖 on_upgrade_prepare: 非 Raid 类型，直接返回
function TestControllerUpgrade:test_on_upgrade_prepare_non_raid()
    local result = self.ctrl_upgrade:on_upgrade_prepare(nil, 'system1', 'Drive', '/path/cfg', nil, {})
    lu.assertNil(result)
end

-- 覆盖 on_upgrade_process: 非 Raid 类型，直接返回
function TestControllerUpgrade:test_on_upgrade_process_non_raid()
    local result = self.ctrl_upgrade:on_upgrade_process(nil, 'system1', 'Drive', '/path/file', {})
    lu.assertNil(result)
end

-- 覆盖 on_upgrade_finish: 非 Raid 类型，直接返回
function TestControllerUpgrade:test_on_upgrade_finish_non_raid()
    local result = self.ctrl_upgrade:on_upgrade_finish(nil, 'system1', 'Drive', {})
    lu.assertNil(result)
end

-- 覆盖 on_upgrade_finish: Raid 类型，清理资源
function TestControllerUpgrade:test_on_upgrade_finish_raid()
    local client = require('storage.client')
    local orig_reply = client.UpdateServiceUpdateServiceFinishReply
    client.UpdateServiceUpdateServiceFinishReply = function() end

    self.ctrl_upgrade.is_controller_upgrading = true
    self.ctrl_upgrade.cfgs = { fake = 1 }
    self.ctrl_upgrade:on_upgrade_finish(nil, 'system1', 'Raid', {})
    lu.assertFalse(self.ctrl_upgrade.is_controller_upgrading)
    lu.assertEquals(next(self.ctrl_upgrade.cfgs), nil)

    client.UpdateServiceUpdateServiceFinishReply = orig_reply
end

TestDriveUpgrade = {}

function TestDriveUpgrade:setUp()
    self.drive_upgrade = require('upgrade.drive_upgrade')
    self._orig_tasks = self.drive_upgrade.tasks
end

function TestDriveUpgrade:tearDown()
    if self._orig_tasks ~= nil then
        self.drive_upgrade.tasks = self._orig_tasks
    end
end

-- 覆盖 on_upgrade_prepare: 非 Drive 类型，直接返回
function TestDriveUpgrade:test_on_upgrade_prepare_non_drive()
    local result = self.drive_upgrade:on_upgrade_prepare(nil, 'system1', 'Raid', '/path/cfg', '/path/hpm', {})
    lu.assertNil(result)
end

-- 覆盖 on_upgrade_process: 非 Drive 类型，直接返回
function TestDriveUpgrade:test_on_upgrade_process_non_drive()
    local result = self.drive_upgrade:on_upgrade_process(nil, 'system1', 'Raid', '/path/file', {})
    lu.assertNil(result)
end

-- 覆盖 on_active_process: 非 Drive 类型，直接返回
function TestDriveUpgrade:test_on_active_process_non_drive()
    local result = self.drive_upgrade:on_active_process(nil, 'system1', 'Raid')
    lu.assertNil(result)
end

-- 覆盖 on_upgrade_finish: 非 Drive 类型，直接返回
function TestDriveUpgrade:test_on_upgrade_finish_non_drive()
    local result = self.drive_upgrade:on_upgrade_finish(nil, 'system1', 'Raid', {})
    lu.assertNil(result)
end