From 5620534ad6ec9a698c16866f00453d334cdd8c98 Mon Sep 17 00:00:00 2001
From: Gerkiz <fs@nvfs.se>
Date: Wed, 4 Oct 2023 23:28:32 +0200
Subject: [PATCH] Health boost module

Delay the removal of given entity so other functions can check if the given entity was previously tracked.
---
 modules/biter_health_booster_v2.lua | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/modules/biter_health_booster_v2.lua b/modules/biter_health_booster_v2.lua
index 83628905..b90c613b 100644
--- a/modules/biter_health_booster_v2.lua
+++ b/modules/biter_health_booster_v2.lua
@@ -8,6 +8,8 @@ local Event = require 'utils.event'
 local LootDrop = require 'modules.mobs_drop_loot'
 local WD = require 'modules.wave_defense.table'
 local Global = require 'utils.global'
+local Task = require 'utils.task'
+local Token = require 'utils.token'
 
 local floor = math.floor
 local insert = table.insert
@@ -88,6 +90,18 @@ local entity_types = {
     ['unit-spawner'] = true
 }
 
+local remove_unit_token =
+    Token.register(
+    function(data)
+        local unit_number = data.unit_number
+        if not unit_number then
+            return
+        end
+
+        this.biter_health_boost_units[unit_number] = nil
+    end
+)
+
 local function loaded_biters(event)
     local cause = event.cause
     local entity = event.entity
@@ -363,7 +377,7 @@ local function on_entity_died(event)
     local wave_count = WD.get_wave()
 
     if health_pool then
-        this.biter_health_boost_units[unit_number] = nil
+        Task.set_timeout_in_ticks(30, remove_unit_token, {unit_number = unit_number})
         if health_pool[3] and health_pool[3].healthbar_id then
             if this.enable_boss_loot then
                 if random(1, 128) == 1 then