diff --git a/pkg/demoinfocs/common/equipment.go b/pkg/demoinfocs/common/equipment.go index d03e989f..d4f0e3a4 100644 --- a/pkg/demoinfocs/common/equipment.go +++ b/pkg/demoinfocs/common/equipment.go @@ -281,7 +281,18 @@ func MapEquipment(eqName string) EquipmentType { wep = EqKnife } else { // If the eqName isn't known it will be EqUnknown as that is the default value for EquipmentType - wep = eqNameToWeapon[eqName] + if strings.HasPrefix(eqName, "m4a1_silencer") { + wep = EqM4A1 + } else if strings.HasPrefix(eqName, "vesthelm") { + wep = EqHelmet + } else { + for name := range eqNameToWeapon { + if strings.HasPrefix(eqName, name) { + wep = eqNameToWeapon[name] + break + } + } + } } return wep diff --git a/pkg/demoinfocs/common/equipment_test.go b/pkg/demoinfocs/common/equipment_test.go index 300c7fe2..34bb6b62 100644 --- a/pkg/demoinfocs/common/equipment_test.go +++ b/pkg/demoinfocs/common/equipment_test.go @@ -24,6 +24,8 @@ func TestMapEquipment(t *testing.T) { assert.Equal(t, EqKnife, MapEquipment("weapon_knife_butterfly"), "'weapon_knife_butterfly' should be mapped to EqKnife") assert.Equal(t, EqM4A4, MapEquipment("weapon_m4a1"), "'weapon_m4a1' should be mapped to EqM4A4") // This is correct, weapon_m4a1 == M4A4 assert.Equal(t, EqM4A1, MapEquipment("weapon_m4a1_silencer"), "'weapon_m4a1_silencer' should be mapped to EqM4A1") + assert.Equal(t, EqKevlar, MapEquipment("weapon_vest"), "'weapon_vest' should be mapped to EqKevlar") + assert.Equal(t, EqHelmet, MapEquipment("weapon_vesthelm"), "'weapon_vesthelm' should be mapped to EqHelmet") assert.Equal(t, EqUnknown, MapEquipment("asdf"), "'asdf' should be mapped to EqUnknown") }