-- 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 file_sec = require 'utils.file'
local utils = require 'mc.utils'
local skynet = require 'skynet'
local base_msg = require 'messages.base'
local enums = require 'macros.power_mgmt_enums'
local capacitor_mgmt = require 'capacitor_mgmt'

TestCapacitorMgmt = {}

function TestCapacitorMgmt:test_calibrate_timers()
    local capacitor_mgmt_mock = {
        system_power_instance = {
            get_calibrate_accumulated_time = function ()
                return 180
            end,
            set_calibrate_accumulated_time = function ()
                return
            end
        },
        last_timers = 0
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.calibrate_timers(capacitor_mgmt_mock, 180)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    ok, rsp = pcall(function ()
        return capacitor_mgmt.calibrate_timers(capacitor_mgmt_mock, 179)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_onepower_alarm_status()
    local capacitor_mgmt_mock = {}

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_alarm_status(capacitor_mgmt_mock,
            {TPSUStatusHigh = 1, TPSUStatusLow = 0})
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_alarm_status(capacitor_mgmt_mock,
            {TPSUStatusHigh = 0, TPSUStatusLow = 1})
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_alarm_status(capacitor_mgmt_mock,
            {TPSUStatusHigh = 0, TPSUStatusLow = 0})
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_onepower_sleep_status()
    local capacitor_mgmt_mock = {}

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_sleep_status(capacitor_mgmt_mock,
            {SleepMode = enums.SLEEP_MODE.DEEP_SLEEP_STR})
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_sleep_status(capacitor_mgmt_mock,
            {SleepMode = enums.SLEEP_MODE.NORMAL_STR})
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_onepower_upgrade_status()
    local capacitor_mgmt_mock = {}

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_upgrade_status(capacitor_mgmt_mock,
            {IsUpgrading = true})
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_onepower_upgrade_status(capacitor_mgmt_mock,
            {IsUpgrading = false})
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_system_averagepower_condition()
    local capacitor_mgmt_mock = {
        power_mgmt_metric = {
            get_system_inputpower_consumption = function ()
                return 301
            end,
            get_system_rated_power = function ()
                return 1000
            end
        },
        window_full_flag = 0,
        avapower_statistics_count = 1
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_averagepower_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.power_mgmt_metric.get_system_inputpower_consumption = function ()
        return 799
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_averagepower_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_system_inputpower_consumption_condition()
    local capacitor_mgmt_mock = {
        power_mgmt_metric = {
            get_system_inputpower_consumption = function ()
                return 301
            end,
            get_system_rated_power = function ()
                return 1000
            end
        }
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_inputpower_consumption_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.power_mgmt_metric.get_system_inputpower_consumption = function ()
        return 799
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_inputpower_consumption_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.power_mgmt_metric.get_system_inputpower_consumption = function ()
        return 300
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_inputpower_consumption_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.power_mgmt_metric.get_system_inputpower_consumption = function ()
        return 800
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_system_inputpower_consumption_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)
end

function TestCapacitorMgmt:test_check_powersupply_requirement_condition_with_system_flag()
    local capacitor_mgmt_mock = {
        psu_service = {
            psu_slot = {{},{}}
        },
        check_onepower_alarm_status = capacitor_mgmt.check_onepower_alarm_status,
        check_onepower_sleep_status = capacitor_mgmt.check_onepower_sleep_status,
        check_onepower_upgrade_status = capacitor_mgmt.check_onepower_upgrade_status
    }
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0
        },
        {
            SleepMode = enums.SLEEP_MODE.DEEP_SLEEP_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        local exit_loop = false
        for _, psu_obj in pairs(psu_objs) do
            acc, exit_loop = func(acc, psu_obj)
            if exit_loop then
                break
            end
        end
        return acc
    end

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_objs[2].SleepMode = enums.SLEEP_MODE.NORMAL_STR
    psu_objs[2].TPSUStatusHigh = 1
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_objs[2].TPSUStatusHigh = 0
    psu_objs[2].TPSUStatusLow = 1
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_objs[2].TPSUStatusLow = 0
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.psu_service.psu_slot = {{},{},{}}
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    c_psu_object.collection.fold = origin_fold_func
end

function TestCapacitorMgmt:test_check_powersupply_requirement_condition_with_execute_flag()
    local capacitor_mgmt_mock = {
        psu_service = {
            psu_slot = {{},{}}
        },
        check_onepower_alarm_status = capacitor_mgmt.check_onepower_alarm_status,
        check_onepower_sleep_status = capacitor_mgmt.check_onepower_sleep_status,
        check_onepower_upgrade_status = capacitor_mgmt.check_onepower_upgrade_status
    }
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false
        },
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        local exit_loop = false
        for _, psu_obj in pairs(psu_objs) do
            acc, exit_loop = func(acc, psu_obj)
            if exit_loop then
                break
            end
        end
        return acc
    end

    psu_objs[2].IsUpgrading = true
    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock, 0)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    psu_objs[2].IsUpgrading = false
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock, 0)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    psu_objs[2].IsUpgrading = true
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock, 1)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_objs[2].IsUpgrading = false
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock, 1)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.psu_service.psu_slot = {{},{},{}}
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_powersupply_requirement_condition(capacitor_mgmt_mock, 0)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    c_psu_object.collection.fold = origin_fold_func
end

function TestCapacitorMgmt:test_capaticors_calibrate_execute()
    local capacitor_mgmt_mock = {
        psu_service = {
            psu_slot = {{},{},{},{}}
        },
        check_calibrate_execute_condition = function ()
            return true
        end,
        capaticor_calibrate_execute = function ()
            return true
        end
    }
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false,
            get_power_capacitor_calibrate_time = function ()
            end
        },
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false,
            get_power_capacitor_calibrate_time = function ()
            end
        },
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false,
            get_power_capacitor_calibrate_time = function ()
            end
        },
        {
            SleepMode = enums.SLEEP_MODE.NORMAL_STR,
            TPSUStatusHigh = 0,
            TPSUStatusLow = 0,
            IsUpgrading = false,
            get_power_capacitor_calibrate_time = function ()
            end
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        local exit_loop = false
        for _, psu_obj in pairs(psu_objs) do
            acc, exit_loop = func(acc, psu_obj)
            if exit_loop then
                break
            end
        end
        return acc
    end

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticors_calibrate_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    capacitor_mgmt_mock.check_calibrate_execute_condition = function ()
        return false
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticors_calibrate_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.check_calibrate_execute_condition = function ()
        return true
    end
    capacitor_mgmt_mock.capaticor_calibrate_execute = function ()
        return false
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticors_calibrate_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.capaticor_calibrate_execute = function ()
        return true
    end
    capacitor_mgmt_mock.psu_service.psu_slot = {{},{},{},{},{}}
    ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticors_calibrate_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.psu_service.psu_slot = {{},{},{},{}}
    ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticors_calibrate_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    c_psu_object.collection.fold = origin_fold_func
end

function TestCapacitorMgmt:test_check_capaticor_calibrate_supported()
    local capacitor_mgmt_mock = {
        calibrate_supported = 0
    }

    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            get_power_capacitor_calibrate = function ()
                return base_msg.ActionNotSupported()
            end
        },
        {
            get_power_capacitor_calibrate = function ()
                return base_msg.ActionNotSupported()
            end
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        local exit_loop = false
        for _, psu_obj in pairs(psu_objs) do
            acc, exit_loop = func(acc, psu_obj)
            if exit_loop then
                break
            end
        end
        return acc
    end

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_capaticor_calibrate_supported(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_objs[1].get_power_capacitor_calibrate = function ()
        return nil
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_capaticor_calibrate_supported(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    psu_objs[2].get_power_capacitor_calibrate = function ()
        return nil
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_capaticor_calibrate_supported(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)

    c_psu_object.collection.fold = origin_fold_func
end

function TestCapacitorMgmt:test_check_calibrate_system_condition()
    local capacitor_mgmt_mock = {
        check_powersupply_requirement_condition = function ()
            return false
        end,
        check_system_averagepower_condition = function ()
            return false
        end
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_system_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.check_powersupply_requirement_condition = function ()
        return true
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_system_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.check_system_averagepower_condition = function ()
        return true
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_system_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_check_calibrate_execute_condition()
    local capacitor_mgmt_mock = {
        check_powersupply_requirement_condition = function ()
            return false
        end,
        check_system_inputpower_consumption_condition = function ()
            return false
        end
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_execute_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.check_powersupply_requirement_condition = function ()
        return true
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_execute_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    capacitor_mgmt_mock.check_system_inputpower_consumption_condition = function ()
        return true
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.check_calibrate_execute_condition(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_capaticor_calibrate_execute()
    local capacitor_mgmt_mock = {}
    local psu_obj = {
        set_power_capacitor_calibrate = function ()
            return base_msg.ActionNotSupported()
        end
    }

    local ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticor_calibrate_execute(capacitor_mgmt_mock, psu_obj)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(rsp, false)

    psu_obj.set_power_capacitor_calibrate = function ()
        return nil
    end
    ok, rsp = pcall(function ()
        return capacitor_mgmt.capaticor_calibrate_execute(capacitor_mgmt_mock, psu_obj)
    end)
    lu.assertTrue(ok)
    lu.assertTrue(rsp)
end

function TestCapacitorMgmt:test_set_life_prediction_status()
    local capacitor_mgmt_mock = {
        life_prediction_status = "Null"
    }

    local ok = pcall(function ()
        capacitor_mgmt.set_life_prediction_status(capacitor_mgmt_mock, "In progress")
    end)
    lu.assertTrue(ok)
    lu.assertEquals(capacitor_mgmt_mock.life_prediction_status, "In progress")

    ok = pcall(function ()
        capacitor_mgmt.set_life_prediction_status(capacitor_mgmt_mock, "In progress")
    end)
    lu.assertTrue(ok)
    lu.assertEquals(capacitor_mgmt_mock.life_prediction_status, "In progress")
end

function TestCapacitorMgmt:test_dump_log()
    local capacitor_mgmt_mock = {
        dod_freq_count_map = {
            [1] = {
                [1] = {
                    [1] = 100,
                    [2] = 200
                },
                [2] = {
                    [1] = 150
                }
            }
        }
    }

    local ok = pcall(function ()
        capacitor_mgmt.dump_log(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)

    capacitor_mgmt_mock.dod_freq_count_map = {}
    ok = pcall(function ()
        capacitor_mgmt.dump_log(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)
end

function TestCapacitorMgmt:test_predicted_lifespan_success()
    local capacitor_mgmt_mock = {
        dod_freq_count_map = {}
    }
    local model_level = 0x568
    local current_dod, current_freq, current_temp = 1, 1, 1
    local psu_obj = {
        ps_id = 1,
        SlotNumber = 1,
        set_life_prediction_count_index = function(self, dod, freq, temp)
            current_dod = dod + 1
            current_freq = freq + 1
            current_temp = temp + 1
            return nil
        end,
        get_life_prediction_count = function(self)
            local index = (current_dod - 1) << 8 | (current_freq - 1) << 4 | (current_temp - 1)
            return nil, index, 10
        end
    }

    local ok, _, _ = pcall(function ()
        return capacitor_mgmt.predicted_lifespan(capacitor_mgmt_mock, model_level, psu_obj)
    end)
    lu.assertTrue(ok)
end

function TestCapacitorMgmt:test_predicted_lifespan_set_index_failed()
    local capacitor_mgmt_mock = {}
    local model_level = 0x568
    local psu_obj = {
        ps_id = 1,
        set_life_prediction_count_index = function(self, dod, freq, temp)
            return "failed"
        end
    }

    local ok, ret = pcall(function ()
        return capacitor_mgmt.predicted_lifespan(capacitor_mgmt_mock, model_level, psu_obj)
    end)
    lu.assertTrue(ok)
    lu.assertNotEquals(ret, nil)
end

function TestCapacitorMgmt:test_predicted_lifespan_get_count_failed()
    local capacitor_mgmt_mock = {}
    local model_level = 0x568
    local psu_obj = {
        ps_id = 1,
        set_life_prediction_count_index = function(self, dod, freq, temp)
            return nil
        end,
        get_life_prediction_count = function(self)
            return "failed"
        end
    }

    local ok, ret = pcall(function ()
        return capacitor_mgmt.predicted_lifespan(capacitor_mgmt_mock, model_level, psu_obj)
    end)
    lu.assertTrue(ok)
    lu.assertNotEquals(ret, nil)
end

function TestCapacitorMgmt:test_predicted_lifespan_index_mismatch()
    local capacitor_mgmt_mock = {}
    local model_level = 0x568
    local psu_obj = {
        ps_id = 1,
        set_life_prediction_count_index = function(self, dod, freq, temp)
            return nil
        end,
        get_life_prediction_count = function(self)
            return nil, 0x999, 10
        end
    }

    local ok, ret = pcall(function ()
        return capacitor_mgmt.predicted_lifespan(capacitor_mgmt_mock, model_level, psu_obj)
    end)
    lu.assertTrue(ok)
    lu.assertEquals(ret, "")
end

function TestCapacitorMgmt:test_life_prediction_execute_level_match()
    local skynet = require 'skynet'
    local origin_sleep = skynet.sleep
    skynet.sleep = function() end
    
    local capacitor_mgmt_mock = {
        dod_freq_count_map = {}
    }
    -- 使用setmetatable确保方法调用正确
    setmetatable(capacitor_mgmt_mock, {
        __index = function(t, k)
            if k == "predicted_lifespan" then
                return capacitor_mgmt.predicted_lifespan
            end
            return capacitor_mgmt[k]
        end
    })
    
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    
    local current_dod, current_freq, current_temp = 0, 0, 0
    local psu_objs = {
        {
            ps_id = 1,
            SlotNumber = 1,
            get_life_prediction_level = function(self)
                return nil, 0x568
            end,
            set_life_prediction_count_index = function(self, dod, freq, temp)
                current_dod = dod
                current_freq = freq
                current_temp = temp
                return nil
            end,
            get_life_prediction_count = function(self)
                local index = current_dod << 8 | current_freq << 4 | current_temp
                return nil, index, 10
            end
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        for _, psu_obj in pairs(psu_objs) do
            acc = func(acc, psu_obj)
        end
        return acc
    end

    local ok, err = pcall(function ()
        capacitor_mgmt.life_prediction_execute(capacitor_mgmt_mock)
    end)
    if not ok then
        print("Error in test_life_prediction_execute_level_match: " .. tostring(err))
    end
    lu.assertTrue(ok, "life_prediction_execute should execute without error")

    c_psu_object.collection.fold = origin_fold_func
    skynet.sleep = origin_sleep
end

function TestCapacitorMgmt:test_life_prediction_execute_level_mismatch()
    local capacitor_mgmt_mock = {}
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            ps_id = 1,
            get_life_prediction_level = function(self)
                return nil, 0x567
            end
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        for _, psu_obj in pairs(psu_objs) do
            acc = func(acc, psu_obj)
        end
        return acc
    end

    local ok = pcall(function ()
        capacitor_mgmt.life_prediction_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)

    c_psu_object.collection.fold = origin_fold_func
end

function TestCapacitorMgmt:test_life_prediction_execute_get_level_failed()
    local capacitor_mgmt_mock = {}
    local c_psu_object = require 'device.psu'
    local origin_fold_func = c_psu_object.collection.fold
    local psu_objs = {
        {
            ps_id = 1,
            get_life_prediction_level = function(self)
                return "failed"
            end
        }
    }
    c_psu_object.collection.fold = function (_, func, acc)
        for _, psu_obj in pairs(psu_objs) do
            acc = func(acc, psu_obj)
        end
        return acc
    end

    local ok = pcall(function ()
        capacitor_mgmt.life_prediction_execute(capacitor_mgmt_mock)
    end)
    lu.assertTrue(ok)

    c_psu_object.collection.fold = origin_fold_func
end