Make power grid check event force APCs off (#17935)

* replace query with enumerator

* remove unnecessary imports

* prevent from toggling APC during event

* all entity query i guess

Content.Server/StationEvents/Components/PowerGridCheckDisabledComponent.cs

@@ -0,0 +1,11 @@
+using Content.Server.StationEvents.Events;
+
+namespace Content.Server.StationEvents.Components;
+
+/// <summary>
+/// Added to APCs to prevent them from turning on until event end
+/// </summary>
+[RegisterComponent, Access(typeof(PowerGridCheckRule))]
+public sealed class PowerGridCheckDisabledComponent : Component
+{
+}

Content.Server/StationEvents/Events/PowerGridCheckRule.cs

@@ -6,9 +6,7 @@ using Robust.Shared.Utility;
 using System.Threading;
 using Content.Server.Power.EntitySystems;
 using Timer = Robust.Shared.Timing.Timer;
-using System.Linq;
 using Content.Server.GameTicking.Rules.Components;
-using Robust.Shared.Random;
 using Content.Server.Station.Components;
 using Content.Server.StationEvents.Components;
 
@@ -19,6 +17,13 @@ namespace Content.Server.StationEvents.Events
     {
         [Dependency] private readonly ApcSystem _apcSystem = default!;
 
+        public override void Initialize()
+        {
+            base.Initialize();
+
+            SubscribeLocalEvent<PowerGridCheckDisabledComponent, ApcToggleMainBreakerAttemptEvent>(OnApcToggleMainBreaker);
+        }
+
         protected override void Started(EntityUid uid, PowerGridCheckRuleComponent component, GameRuleComponent gameRule, GameRuleStartedEvent args)
         {
             base.Started(uid, component, gameRule, args);
@@ -26,10 +31,11 @@ namespace Content.Server.StationEvents.Events
             if (!TryGetRandomStation(out var chosenStation))
                 return;
 
-            foreach (var (apc, transform) in EntityQuery<ApcComponent, TransformComponent>(true))
+            var query = AllEntityQuery<ApcComponent, TransformComponent>();
+            while (query.MoveNext(out var target, out var apc, out var transform))
             {
                 if (apc.MainBreakerEnabled && CompOrNull<StationMemberComponent>(transform.GridUid)?.Station == chosenStation)
-                    component.Powered.Add(apc.Owner);
+                    component.Powered.Add(target);
             }
 
             RobustRandom.Shuffle(component.Powered);
@@ -51,6 +57,7 @@ namespace Content.Server.StationEvents.Events
                     if (!apcComponent.MainBreakerEnabled)
                         _apcSystem.ApcToggleBreaker(entity, apcComponent);
                 }
+                RemComp<PowerGridCheckDisabledComponent>(entity);
             }
 
             // Can't use the default EndAudio
@@ -88,8 +95,14 @@ namespace Content.Server.StationEvents.Events
                     if (apcComponent.MainBreakerEnabled)
                         _apcSystem.ApcToggleBreaker(selected, apcComponent);
                 }
+                EnsureComp<PowerGridCheckDisabledComponent>(selected);
                 component.Unpowered.Add(selected);
             }
         }
+
+        private void OnApcToggleMainBreaker(EntityUid uid, PowerGridCheckDisabledComponent component, ref ApcToggleMainBreakerAttemptEvent args)
+        {
+            args.Cancelled = true;
+        }
     }
 }