From 236da1cd27a0ff5569113e30782a5f771f254037 Mon Sep 17 00:00:00 2001
From: Tayrtahn <tayrtahn@gmail.com>
Date: Thu, 2 May 2024 08:16:16 -0400
Subject: [PATCH] Make UseDelay SetLength do EnsureComp (#27601)

Make UseDelay.SetLength do EnsureComp
---
 Content.Server/Fluids/EntitySystems/SpraySystem.cs  |  2 +-
 .../Storage/EntitySystems/SharedStorageSystem.cs    |  7 ++-----
 Content.Shared/Timing/UseDelaySystem.cs             | 13 +++++++++----
 3 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/Content.Server/Fluids/EntitySystems/SpraySystem.cs b/Content.Server/Fluids/EntitySystems/SpraySystem.cs
index 40f19aff2b..5499070738 100644
--- a/Content.Server/Fluids/EntitySystems/SpraySystem.cs
+++ b/Content.Server/Fluids/EntitySystems/SpraySystem.cs
@@ -144,7 +144,7 @@ public sealed class SpraySystem : EntitySystem
 
         _audio.PlayPvs(entity.Comp.SpraySound, entity, entity.Comp.SpraySound.Params.WithVariation(0.125f));
 
-        _useDelay.SetLength((entity, useDelay), TimeSpan.FromSeconds(cooldownTime));
+        _useDelay.SetLength(entity.Owner, TimeSpan.FromSeconds(cooldownTime));
         _useDelay.TryResetDelay((entity, useDelay));
     }
 }
diff --git a/Content.Shared/Storage/EntitySystems/SharedStorageSystem.cs b/Content.Shared/Storage/EntitySystems/SharedStorageSystem.cs
index 4670f2dbf8..771513a0a8 100644
--- a/Content.Shared/Storage/EntitySystems/SharedStorageSystem.cs
+++ b/Content.Shared/Storage/EntitySystems/SharedStorageSystem.cs
@@ -135,11 +135,8 @@ public abstract class SharedStorageSystem : EntitySystem
 
     private void OnMapInit(Entity<StorageComponent> entity, ref MapInitEvent args)
     {
-        if (TryComp<UseDelayComponent>(entity, out var useDelayComp))
-        {
-            UseDelay.SetLength((entity, useDelayComp), entity.Comp.QuickInsertCooldown, QuickInsertUseDelayID);
-            UseDelay.SetLength((entity, useDelayComp), entity.Comp.OpenUiCooldown, OpenUiUseDelayID);
-        }
+        UseDelay.SetLength(entity.Owner, entity.Comp.QuickInsertCooldown, QuickInsertUseDelayID);
+        UseDelay.SetLength(entity.Owner, entity.Comp.OpenUiCooldown, OpenUiUseDelayID);
     }
 
     private void OnStorageGetState(EntityUid uid, StorageComponent component, ref ComponentGetState args)
diff --git a/Content.Shared/Timing/UseDelaySystem.cs b/Content.Shared/Timing/UseDelaySystem.cs
index bc2a709175..9816d0185a 100644
--- a/Content.Shared/Timing/UseDelaySystem.cs
+++ b/Content.Shared/Timing/UseDelaySystem.cs
@@ -47,7 +47,7 @@ public sealed class UseDelaySystem : EntitySystem
     {
         // Set default delay length from the prototype
         // This makes it easier for simple use cases that only need a single delay
-        SetLength(ent, ent.Comp.Delay, DefaultId);
+        SetLength((ent, ent.Comp), ent.Comp.Delay, DefaultId);
     }
 
     private void OnUnpaused(Entity<UseDelayComponent> ent, ref EntityUnpausedEvent args)
@@ -62,9 +62,14 @@ public sealed class UseDelaySystem : EntitySystem
     /// <summary>
     /// Sets the length of the delay with the specified ID.
     /// </summary>
-    public bool SetLength(Entity<UseDelayComponent> ent, TimeSpan length, string id = DefaultId)
+    /// <remarks>
+    /// This will add a UseDelay component to the entity if it doesn't have one.
+    /// </remarks>
+    public bool SetLength(Entity<UseDelayComponent?> ent, TimeSpan length, string id = DefaultId)
     {
-        if (ent.Comp.Delays.TryGetValue(id, out var entry))
+        EnsureComp<UseDelayComponent>(ent.Owner, out var comp);
+
+        if (comp.Delays.TryGetValue(id, out var entry))
         {
             if (entry.Length == length)
                 return true;
@@ -73,7 +78,7 @@ public sealed class UseDelaySystem : EntitySystem
         }
         else
         {
-            ent.Comp.Delays.Add(id, new UseDelayInfo(length));
+            comp.Delays.Add(id, new UseDelayInfo(length));
         }
 
         Dirty(ent);
-- 
2.52.0