-- 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 def = require 'common_def'
local method_misc = require 'method_misc'

TestHandlerPassThroughDriveStub = {}

function TestHandlerPassThroughDriveStub:setUp()
    package.loaded['bma.handles.handler_pass_through_drive'] = nil
end

function TestHandlerPassThroughDriveStub:tearDown()
    package.loaded['bma.handles.handler_pass_through_drive'] = nil
end

-- get_val_or_default
function TestHandlerPassThroughDriveStub:test_get_val_or_default()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    lu.assertEquals(c:get_val_or_default(nil, 'd'), 'd')
    lu.assertEquals(c:get_val_or_default('v', 'd'), 'v')
end

-- get_key_by_val
function TestHandlerPassThroughDriveStub:test_get_key_by_val()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local t = { a = 1, b = 2 }
    lu.assertEquals(c:get_key_by_val(t, 2, nil), 'b')
    lu.assertEquals(c:get_key_by_val(t, 9, 'def'), 'def')
end

-- get_written_info: 无匹配返回 INVALID_DWORD
function TestHandlerPassThroughDriveStub:test_get_written_info_no_match()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local list = { { ['ID#'] = 1, RAW_VALUE = 100 } }
    lu.assertEquals(c:get_written_info(99, list), def.STORAGE_INFO_INVALID_DWORD)
    lu.assertEquals(c:get_written_info(1, list), 100)
end

-- check_hw_defined_smart
function TestHandlerPassThroughDriveStub:test_check_hw_defined_smart()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    lu.assertFalse(c:check_hw_defined_smart({
        Oem = { Huawei = { HuaweiSpecificLogPages = nil } },
        SmartDetails = { HuaweiDefinedAttributeItemList = {} }
    }))
    lu.assertTrue(c:check_hw_defined_smart({
        Oem = { Huawei = { HuaweiSpecificLogPages = { 1 } } },
        SmartDetails = { HuaweiDefinedAttributeItemList = { 1 } }
    }))
end

-- get_smart_value_by_attr_id
function TestHandlerPassThroughDriveStub:test_get_smart_value_by_attr_id()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local list = { { ['ID#'] = 5, VALUE = 80 } }
    lu.assertEquals(c:get_smart_value_by_attr_id(5, list), 80)
    lu.assertEquals(c:get_smart_value_by_attr_id(6, list), def.INVALID_U8)
end

-- check_manufacture_statisfied_top3
function TestHandlerPassThroughDriveStub:test_check_manufacture_statisfied_top3()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    if smart_def and smart_def.PD_VENDOR_NAME_INTEL then
        lu.assertTrue(c:check_manufacture_statisfied_top3({ Manufacturer = smart_def.PD_VENDOR_NAME_INTEL }))
    end
    lu.assertFalse(c:check_manufacture_statisfied_top3({ Manufacturer = 'Other' }))
end

-- check_collect_est_lifespan_time: 首次 last==0
function TestHandlerPassThroughDriveStub:test_check_collect_est_lifespan_time_first()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local obj = { est_lifespan_last_up_time = 0, write_amplification_info = {} }
    lu.assertTrue(c:check_collect_est_lifespan_time(obj))
    lu.assertNotNil(obj.est_lifespan_last_up_time)
    lu.assertEquals(obj.write_amplification_info.first_start_flag, 1)
end

-- get_wearout_raw_by_attr_id / get_smart_raw_by_attr_id 覆盖
function TestHandlerPassThroughDriveStub:test_get_wearout_and_smart_raw_by_attr_id()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local list = { { ['ID#'] = 10, RAW_VALUE = 200 } }
    lu.assertEquals(c:get_wearout_raw_by_attr_id(10, list), 200)
    lu.assertEquals(c:get_wearout_raw_by_attr_id(11, list), def.INVALID_U8)
    lu.assertEquals(c:get_smart_raw_by_attr_id(10, list), 200)
    lu.assertEquals(c:get_smart_raw_by_attr_id(11, list), def.INVALID_U32)
end

function TestHandlerPassThroughDriveStub:test_get_pd_wearout_by_hw_defined_min()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    local list = {
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_WEAROUT, RAW_VALUE = 20 },
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_WEAROUT, RAW_VALUE = 10 }
    }
    lu.assertEquals(c:get_pd_wearout_by_hw_defined(list), 10)
end

function TestHandlerPassThroughDriveStub:test_get_pd_wearout_info_sata_hw_defined()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    local obj = {
        Protocol = def.DRIVE_PROTOCOL.SATA,
        Manufacturer = 'ATA',
        estimated_remaining_lifespan_info = { is_support_hw_defined = 1 }
    }
    local data = {
        AttributeItemList = { { ['ID#'] = 1, VALUE = 1 } },
        HuaweiDefinedAttributeItemList = {
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_WEAROUT, RAW_VALUE = 7 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_WEAROUT, RAW_VALUE = 9 }
        }
    }
    lu.assertEquals(c:get_pd_wearout_info(obj, data), 7)
end

function TestHandlerPassThroughDriveStub:test_get_pd_wearout_info_sata_vendor_attr()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    local mfg, attr_id = next(smart_def.VENDOR_ATTR_ID or {})
    if not mfg then
        lu.assertTrue(true)
        return
    end
    local obj = {
        Protocol = def.DRIVE_PROTOCOL.SATA,
        Manufacturer = mfg,
        estimated_remaining_lifespan_info = { is_support_hw_defined = 0 }
    }
    local data = {
        AttributeItemList = { { ['ID#'] = attr_id, VALUE = 88 } }
    }
    lu.assertEquals(c:get_pd_wearout_info(obj, data), 88)
end

function TestHandlerPassThroughDriveStub:test_get_pd_wearout_info_sas_indicator()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    local obj = { Protocol = def.DRIVE_PROTOCOL.SAS }
    local data = { SSMediaUsedEnduranceIndicator = 20 }
    lu.assertEquals(c:get_pd_wearout_info(obj, data), smart_def.ONE_HUNDRED_PERCENT - 20)
end

function TestHandlerPassThroughDriveStub:test_update_hw_smart_info_avg_parse()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF
    local obj = {
        PowerOnHours = 123,
        estimated_remaining_lifespan_info = {}
    }
    local list = {
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_ERASE_COUNT, RAW_VALUE = 'avg:100' },
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_ERASE_COUNT, RAW_VALUE = 'avg:200' },
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_PE_CYCLE, RAW_VALUE = 111 },
        { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_PE_CYCLE, RAW_VALUE = 222 }
    }
    c:update_hw_smart_info(obj, list)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.slc_avg_ec, 100)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.tlc_avg_ec, 200)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.slc_pe_cycle, 111)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.tlc_pe_cycle, 222)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.slc_poh, 123)
end

function TestHandlerPassThroughDriveStub:test_update_est_lifespan_hw_defined_and_fallback()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF

    local obj1 = {
        Name = 'd1',
        MediaType = 0,
        Manufacturer = smart_def.PD_VENDOR_NAME_HUAWEI or 'Huawei',
        SerialNumber = 'sn',
        Protocol = def.DRIVE_PROTOCOL.SATA,
        PowerOnHours = 1,
        PredictedMediaLifeLeftPercent = 77,
        estimated_remaining_lifespan_info = { is_support_hw_defined = 1 }
    }
    c:update_est_lifespan(obj1, {
        HuaweiDefinedAttributeItemList = {
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_ERASE_COUNT, RAW_VALUE = 'avg:1' },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_ERASE_COUNT, RAW_VALUE = 'avg:2' },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_PE_CYCLE, RAW_VALUE = 3 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_PE_CYCLE, RAW_VALUE = 4 }
        }
    })
    lu.assertEquals(obj1.estimated_remaining_lifespan_info.drive_name, 'd1')
    lu.assertEquals(obj1.estimated_remaining_lifespan_info.slc_avg_ec, 1)

    local obj2 = {
        Name = 'd2',
        MediaType = 0,
        Manufacturer = 'Other',
        SerialNumber = 'sn2',
        Protocol = def.DRIVE_PROTOCOL.SATA,
        PowerOnHours = 9,
        PredictedMediaLifeLeftPercent = 66,
        estimated_remaining_lifespan_info = { is_support_hw_defined = 0 }
    }
    c:update_est_lifespan(obj2, {})
    lu.assertEquals(obj2.estimated_remaining_lifespan_info.remn_wearout, 66)
    lu.assertEquals(obj2.estimated_remaining_lifespan_info.power_on_hours, 9)
end

function TestHandlerPassThroughDriveStub:test_update_write_amplification_vendor_hw_defined_and_sas()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF

    local mfg_nand, nand_id = next(smart_def.NAND_WRITTEN_ATTR_ID or {})
    local _, host_id = next(smart_def.HOST_WRITTEN_ATTR_ID or {})
    if mfg_nand and nand_id and host_id then
        local obj = { Manufacturer = mfg_nand, write_amplification_info = {} }
        c:update_sata_write_amplification_by_vendor(obj, {
            AttributeItemList = {
                { ['ID#'] = nand_id, RAW_VALUE = 10 },
                { ['ID#'] = host_id, RAW_VALUE = 20 }
            }
        })
        lu.assertEquals(obj.write_amplification_info.vendor_valid_flag, 1)
    end

    local obj2 = { write_amplification_info = {} }
    c:update_sata_write_amplification_by_hw_defined(obj2, {
        HuaweiDefinedAttributeItemList = {
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_NAND_WRITTEN_L, RAW_VALUE = 1 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_NAND_WRITTEN_H, RAW_VALUE = 2 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_HOST_WRITTEN_L, RAW_VALUE = 3 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_HOST_WRITTEN_H, RAW_VALUE = 4 }
        }
    })
    lu.assertEquals(obj2.write_amplification_info.hw_defined_host_write_h, 4)

    local obj3 = { write_amplification_info = {} }
    c:update_sas_write_amplification(obj3, { NandWrittenMiB = nil, HostWrittenMiB = nil })
    lu.assertNil(obj3.write_amplification_info.vendor_valid_flag)
    c:update_sas_write_amplification(obj3, { NandWrittenMiB = 9, HostWrittenMiB = nil })
    lu.assertEquals(obj3.write_amplification_info.vendor_nand_write, 9)
    lu.assertEquals(obj3.write_amplification_info.vendor_valid_flag, 1)
end

function TestHandlerPassThroughDriveStub:test_check_collect_est_lifespan_time_interval_and_no_update()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local now = math.floor(os.time())
    local obj = { est_lifespan_last_up_time = now, write_amplification_info = {} }
    lu.assertFalse(c:check_collect_est_lifespan_time(obj))

    local obj2 = { est_lifespan_last_up_time = now - (3600 * 24) - 1, write_amplification_info = {} }
    lu.assertTrue(c:check_collect_est_lifespan_time(obj2))
    lu.assertTrue(obj2.est_lifespan_last_up_time >= now - 2)
end

function TestHandlerPassThroughDriveStub:test_update_sata_write_amplification_choose_branch()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local hw_called, vendor_called = false, false
    c.update_sata_write_amplification_by_hw_defined = function() hw_called = true end
    c.update_sata_write_amplification_by_vendor = function() vendor_called = true end
    c:update_sata_write_amplification({ estimated_remaining_lifespan_info = { is_support_hw_defined = 1 } }, {})
    lu.assertTrue(hw_called)
    c:update_sata_write_amplification({ estimated_remaining_lifespan_info = { is_support_hw_defined = 0 } }, {})
    lu.assertTrue(vendor_called)
end

function TestHandlerPassThroughDriveStub:test_update_smart_info_sata_calls_est_lifespan_and_write_amp()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local smart_def = def.PD_SMART_DEF

    local est_called, amp_called = false, false
    c.update_est_lifespan = function() est_called = true end
    c.update_sata_write_amplification = function() amp_called = true end

    local obj = {
        Protocol = def.DRIVE_PROTOCOL.SATA,
        MediaType = 0,
        Manufacturer = smart_def.PD_VENDOR_NAME_HUAWEI or 'Huawei',
        PowerOnHours = (30 * 24),
        write_amplification_info = {},
        estimated_remaining_lifespan_info = { is_support_hw_defined = 0 }
    }
    local data = {
        Oem = { Huawei = { HuaweiSpecificLogPages = { 1 } } },
        SmartDetails = {
            HuaweiDefinedAttributeItemList = { 1 },
        },
        HuaweiDefinedAttributeItemList = {
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_SLC_WEAROUT, RAW_VALUE = 1 },
            { ['ID#'] = smart_def.HW_DEFINED_SMART_ATTRIBUTE_ID_TLC_WEAROUT, RAW_VALUE = 2 }
        }
    }
    c:update_smart_info(obj, data)
    lu.assertEquals(obj.estimated_remaining_lifespan_info.is_support_hw_defined, 1)
    lu.assertEquals(obj.write_amplification_info.hw_defined_valid_flag, 1)
    lu.assertEquals(obj.write_amplification_info.update_support_flag, 1)
    lu.assertTrue(est_called)
    lu.assertTrue(amp_called)
end

function TestHandlerPassThroughDriveStub:test_update_hw_info_and_drive_manufacturer_paths()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()
    local obj = { Manufacturer = 'ATA' }
    -- 使用 method_misc:get_manufacturer_name_by_oui 的真实映射（覆盖 update_hw_info 的 ATA->OUI 分支）
    c:update_hw_info(obj, { FirmwareVersion = '1', VendorID = '16', Temperature = 33, PowerOnHours = 2, PowerState = 1, ManufacturerOUI = 'A428B7' })
    lu.assertEquals(obj.Revision, '1')
    lu.assertEquals(obj.ManufacturerId, 16)
    lu.assertEquals(obj.TemperatureCelsius, 33)
    lu.assertEquals(obj.PowerOnHours, 2)
    lu.assertEquals(obj.PowerState, 1)
    lu.assertEquals(obj.Manufacturer, 'YANGTZE MEMORY')

    local vid = next(def.VENDOR_ID_2_VENDOR_NAME or {})
    if vid then
        local obj2 = {}
        c:update_drive_manufacturer(obj2, { Oem = { Huawei = { ManufacturerOUI = string.format('%x', vid) } } })
        lu.assertEquals(obj2.Manufacturer, def.VENDOR_ID_2_VENDOR_NAME[vid])
    end
end

function TestHandlerPassThroughDriveStub:test_update_drive_info_and_update_drive_and_dump_task()
    local handler = require 'bma.handles.handler_pass_through_drive'
    local c = handler.new()

    local obj = {
        Name = 'd',
        write_amplification_info = {},
        estimated_remaining_lifespan_info = { is_support_hw_defined = 0 },
        new_task = function()
            return {
                loop = function(self, fn)
                    fn({ is_exit = false })
                    return { set_timeout_ms = function() end }
                end
            }
        end,
    }
    obj.update_asset_data_info = function(self) self.asset_updated = true end
    obj.record_write_amp_log = function(self) self.logged = true end
    c:update_drive_info(obj, {
        Manufacturer = 'ATA',
        Model = 'M',
        Protocol = 'SATA',
        MediaType = 'SSD',
        SerialNumber = 'SN',
        CapacityBytes = 1024 * 1024,
        RotationSpeedRPM = 0,
        CapableSpeedGbs = 6.0,
        NegotiatedSpeedGbs = 3.0,
        MediaErrorCount = 1,
        OtherErrorCount = 2,
        PredFailCount = 3,
        Status = 'Active',
        Oem = { Huawei = { FirmwareVersion = 'F', VendorID = '1', Temperature = 1, PowerOnHours = 1, PowerState = 1 } },
        SmartDetails = { AttributeItemList = {} }
    })
    lu.assertEquals(obj.Model, 'M')
    lu.assertEquals(obj.SerialNumber, 'SN')

    c:update_drive(obj, { FirmwareVersion = 'FW', Oem = { Huawei = { SilkScreen = 'x' } } })
    lu.assertTrue(obj.asset_updated)

    -- start_dump_task: early return + normal path
    obj.logged = false
    obj.ObjectName = 'o'
    obj.write_amplification_info.vendor_valid_flag = 0
    obj.write_amplification_info.hw_defined_valid_flag = 0
    c:start_dump_task(obj, {})
    lu.assertFalse(obj.logged)
    obj.write_amplification_info.vendor_valid_flag = 1
    c.check_collect_est_lifespan_time = function() return true end
    c:start_dump_task(obj, {})
    lu.assertTrue(obj.logged)
end
