Disease fixes (#530)

Rxup · Mar 23, 2024 · 3adfd00 · 3adfd00
1 parent 8ba1fcf
commit 3adfd00
Show file tree

Hide file tree

Showing 21 changed files with 339 additions and 156 deletions.
diff --git a/Content.IntegrationTests/Tests/Backmen/Disease/TryAddDisease.cs b/Content.IntegrationTests/Tests/Backmen/Disease/TryAddDisease.cs
@@ -37,10 +37,15 @@ await server.WaitAssertion(() =>
         {
             await server.WaitAssertion(() =>
             {
-                diseaseSystem.TryAddDisease(sickEntity, diseaseProto);
+                diseaseSystem.TryAddDisease(sickEntity, diseaseProto.ID);
             });
             await server.WaitIdleAsync();
-            server.RunTicks(1);
+            server.RunTicks(5);
+            await server.WaitAssertion(() =>
+            {
+                if(!entManager.HasComponent<DiseasedComponent>(sickEntity))
+                    Assert.Fail("MobHuman has not DiseasedComponent");
+            });
             if (!entManager.TryGetComponent<DiseaseCarrierComponent>(sickEntity, out var diseaseCarrierComponent))
             {
                 Assert.Fail("MobHuman has not DiseaseCarrierComponent");

diff --git a/Content.Server/Backmen/Disease/Components/DiseaseInfectionSpeadEvent.cs b/Content.Server/Backmen/Disease/Components/DiseaseInfectionSpeadEvent.cs
@@ -0,0 +1,10 @@
+using Content.Shared.Backmen.Disease;
+
+namespace Content.Server.Backmen.Disease.Components;
+
+public sealed class DiseaseInfectionSpreadEvent : EntityEventArgs
+{
+    public EntityUid Owner { get; init; } = default!;
+    public DiseasePrototype Disease { get; init; } = default!;
+    public float Range { get; init; } = default!;
+}
diff --git a/Content.Server/Backmen/Disease/Cures/DiseaseBedrestCure.cs b/Content.Server/Backmen/Disease/Cures/DiseaseBedrestCure.cs
@@ -2,6 +2,7 @@
 using Content.Shared.Backmen.Disease;
 using Content.Shared.Bed.Sleep;
 using Content.Shared.Buckle.Components;
+using Robust.Shared.Prototypes;
 
 namespace Content.Server.Backmen.Disease.Cures;
 
@@ -22,23 +23,33 @@ public override string CureText()
     {
         return (Loc.GetString("diagnoser-cure-bedrest", ("time", MaxLength), ("sleep", (MaxLength / SleepMultiplier))));
     }
+
+    public override object GenerateEvent(Entity<DiseaseCarrierComponent> ent, ProtoId<DiseasePrototype> disease)
+    {
+        return new DiseaseCureArgs<DiseaseBedrestCure>(ent, disease, this);
+    }
 }
 
 public sealed partial class DiseaseCureSystem
 {
-    private void DiseaseBedrestCure(DiseaseCureArgs args, DiseaseBedrestCure ds)
+    private void DiseaseBedrestCure(Entity<DiseaseCarrierComponent> ent, ref DiseaseCureArgs<DiseaseBedrestCure> args)
     {
-        if (!TryComp<BuckleComponent>(args.DiseasedEntity, out var buckle) ||
-            !HasComp<HealOnBuckleComponent>(buckle.BuckledTo))
+        if(args.Handled)
+            return;
+
+        args.Handled = true;
+
+        if (!_buckleQuery.TryGetComponent(args.DiseasedEntity, out var buckle) ||
+            !_healOnBuckleQuery.HasComponent(buckle.BuckledTo))
             return;
 
         var ticks = 1;
-        if (HasComp<SleepingComponent>(args.DiseasedEntity))
-            ticks *= ds.SleepMultiplier;
+        if (_sleepingComponentQuery.HasComponent(args.DiseasedEntity))
+            ticks *= args.DiseaseCure.SleepMultiplier;
 
         if (buckle.Buckled)
-            ds.Ticker += ticks;
-        if (ds.Ticker >= ds.MaxLength)
+            args.DiseaseCure.Ticker += ticks;
+        if (args.DiseaseCure.Ticker >= args.DiseaseCure.MaxLength)
         {
             _disease.CureDisease(args.DiseasedEntity, args.Disease);
         }

diff --git a/Content.Server/Backmen/Disease/Cures/DiseaseBodyTemperatureCure.cs b/Content.Server/Backmen/Disease/Cures/DiseaseBodyTemperatureCure.cs
@@ -1,5 +1,6 @@
 using Content.Server.Temperature.Components;
 using Content.Shared.Backmen.Disease;
+using Robust.Shared.Prototypes;
 
 namespace Content.Server.Backmen.Disease.Cures;
 
@@ -20,16 +21,26 @@ public override string CureText()
 
         return Loc.GetString("diagnoser-cure-temp-both", ("max", Math.Round(Max)), ("min", Math.Round(Min)));
     }
+
+    public override object GenerateEvent(Entity<DiseaseCarrierComponent> ent, ProtoId<DiseasePrototype> disease)
+    {
+        return new DiseaseCureArgs<DiseaseBodyTemperatureCure>(ent, disease, this);
+    }
 }
 
 public sealed partial class DiseaseCureSystem
 {
-    private void DiseaseBodyTemperatureCure(DiseaseCureArgs args, DiseaseBodyTemperatureCure ds)
+    private void DiseaseBodyTemperatureCure(Entity<DiseaseCarrierComponent> ent, ref DiseaseCureArgs<DiseaseBodyTemperatureCure> args)
     {
-        if (!TryComp<TemperatureComponent>(args.DiseasedEntity, out var temp))
+        if(args.Handled)
+            return;
+
+        args.Handled = true;
+
+        if (!_temperatureQuery.TryGetComponent(args.DiseasedEntity, out var temp))
             return;
 
-        if(temp.CurrentTemperature > ds.Min && temp.CurrentTemperature < float.MaxValue)
+        if(temp.CurrentTemperature > args.DiseaseCure.Min && temp.CurrentTemperature < float.MaxValue)
         {
             _disease.CureDisease(args.DiseasedEntity, args.Disease);
         }

diff --git a/Content.Server/Backmen/Disease/Cures/DiseaseCureSystem.cs b/Content.Server/Backmen/Disease/Cures/DiseaseCureSystem.cs
@@ -1,40 +1,34 @@
-using Content.Shared.Backmen.Disease;
+using Content.Server.Bed.Components;
+using Content.Server.Body.Components;
+using Content.Server.Temperature.Components;
+using Content.Shared.Backmen.Disease;
+using Content.Shared.Bed.Sleep;
+using Content.Shared.Buckle.Components;
 
 namespace Content.Server.Backmen.Disease.Cures;
 
 public sealed partial class DiseaseCureSystem : EntitySystem
 {
     [Dependency] private readonly DiseaseSystem _disease = default!;
+    private EntityQuery<BuckleComponent> _buckleQuery;
+    private EntityQuery<HealOnBuckleComponent> _healOnBuckleQuery;
+    private EntityQuery<SleepingComponent> _sleepingComponentQuery;
+    private EntityQuery<BloodstreamComponent> _bloodstreamQuery;
+    private EntityQuery<TemperatureComponent> _temperatureQuery;
+
     public override void Initialize()
     {
         base.Initialize();
 
-        SubscribeLocalEvent<DiseaseCureArgs>(DoGenericCure);
-    }
-
-    private void DoGenericCure(DiseaseCureArgs args)
-    {
-        if(args.Handled)
-            return;
+        SubscribeLocalEvent<DiseaseCarrierComponent, DiseaseCureArgs<DiseaseBedrestCure>>(DiseaseBedrestCure);
+        SubscribeLocalEvent<DiseaseCarrierComponent, DiseaseCureArgs<DiseaseBodyTemperatureCure>>(DiseaseBodyTemperatureCure);
+        SubscribeLocalEvent<DiseaseCarrierComponent, DiseaseCureArgs<DiseaseJustWaitCure>>(DiseaseJustWaitCure);
+        SubscribeLocalEvent<DiseaseCarrierComponent, DiseaseCureArgs<DiseaseReagentCure>>(DiseaseReagentCure);
 
-        switch (args.DiseaseCure)
-        {
-            case DiseaseBedrestCure ds1:
-                args.Handled = true;
-                DiseaseBedrestCure(args, ds1);
-                return;
-            case DiseaseBodyTemperatureCure ds2:
-                args.Handled = true;
-                DiseaseBodyTemperatureCure(args, ds2);
-                return;
-            case DiseaseJustWaitCure ds3:
-                args.Handled = true;
-                DiseaseJustWaitCure(args, ds3);
-                return;
-            case DiseaseReagentCure ds4:
-                args.Handled = true;
-                DiseaseReagentCure(args, ds4);
-                return;
-        }
+        _buckleQuery = GetEntityQuery<BuckleComponent>();
+        _healOnBuckleQuery = GetEntityQuery<HealOnBuckleComponent>();
+        _sleepingComponentQuery = GetEntityQuery<SleepingComponent>();
+        _bloodstreamQuery = GetEntityQuery<BloodstreamComponent>();
+        _temperatureQuery = GetEntityQuery<TemperatureComponent>();
     }
 }
diff --git a/Content.Server/Backmen/Disease/Cures/DiseaseJustWaitCure.cs b/Content.Server/Backmen/Disease/Cures/DiseaseJustWaitCure.cs
@@ -1,4 +1,5 @@
 using Content.Shared.Backmen.Disease;
+using Robust.Shared.Prototypes;
 
 namespace Content.Server.Backmen.Disease.Cures;
 
@@ -21,14 +22,24 @@ public override string CureText()
     {
         return Loc.GetString("diagnoser-cure-wait", ("time", MaxLength));
     }
+
+    public override object GenerateEvent(Entity<DiseaseCarrierComponent> ent, ProtoId<DiseasePrototype> disease)
+    {
+        return new DiseaseCureArgs<DiseaseJustWaitCure>(ent, disease, this);
+    }
 }
 
 public sealed partial class DiseaseCureSystem
 {
-    private void DiseaseJustWaitCure(DiseaseCureArgs args, DiseaseJustWaitCure ds)
+    private void DiseaseJustWaitCure(Entity<DiseaseCarrierComponent> ent, ref DiseaseCureArgs<DiseaseJustWaitCure> args)
     {
-        ds.Ticker++;
-        if (ds.Ticker >= ds.MaxLength)
+        if(args.Handled)
+            return;
+
+        args.Handled = true;
+
+        args.DiseaseCure.Ticker++;
+        if (args.DiseaseCure.Ticker >= args.DiseaseCure.MaxLength)
         {
             _disease.CureDisease(args.DiseasedEntity, args.Disease);
         }

diff --git a/Content.Server/Backmen/Disease/Cures/DiseaseReagentCure.cs b/Content.Server/Backmen/Disease/Cures/DiseaseReagentCure.cs
@@ -24,25 +24,35 @@ public override string CureText()
             return string.Empty;
         return (Loc.GetString("diagnoser-cure-reagent", ("units", Min), ("reagent", reagentProt.LocalizedName)));
     }
+
+    public override object GenerateEvent(Entity<DiseaseCarrierComponent> ent, ProtoId<DiseasePrototype> disease)
+    {
+        return new DiseaseCureArgs<DiseaseReagentCure>(ent, disease, this);
+    }
 }
 
 public sealed partial class DiseaseCureSystem
 {
-    private void DiseaseReagentCure(DiseaseCureArgs args, DiseaseReagentCure ds)
+    private void DiseaseReagentCure(Entity<DiseaseCarrierComponent> ent, ref DiseaseCureArgs<DiseaseReagentCure> args)
     {
-        if (!TryComp<BloodstreamComponent>(args.DiseasedEntity, out var bloodstream)
+        if(args.Handled)
+            return;
+
+        args.Handled = true;
+
+        if (!_bloodstreamQuery.TryGetComponent(args.DiseasedEntity, out var bloodstream)
             || bloodstream.ChemicalSolution == null)
             return;
 
         var chemicalSolution = bloodstream.ChemicalSolution.Value;
 
         var quant = FixedPoint2.Zero;
-        if (ds.Reagent != null && chemicalSolution.Comp.Solution.ContainsReagent(ds.Reagent.Value))
+        if (args.DiseaseCure.Reagent != null && chemicalSolution.Comp.Solution.ContainsReagent(args.DiseaseCure.Reagent.Value))
         {
-            quant = chemicalSolution.Comp.Solution.GetReagentQuantity(ds.Reagent.Value);
+            quant = chemicalSolution.Comp.Solution.GetReagentQuantity(args.DiseaseCure.Reagent.Value);
         }
 
-        if (quant >= ds.Min)
+        if (quant >= args.DiseaseCure.Min)
         {
             _disease.CureDisease(args.DiseasedEntity, args.Disease);
         }

diff --git a/Content.Server/Backmen/Disease/DiseaseInfectionSpread.cs b/Content.Server/Backmen/Disease/DiseaseInfectionSpread.cs
@@ -0,0 +1,32 @@
+using System.Threading;
+using System.Threading.Tasks;
+using Content.Server.Backmen.Disease.Components;
+using Content.Server.Backmen.Disease.Effects;
+using Robust.Shared.CPUJob.JobQueues;
+using Robust.Shared.Timing;
+
+namespace Content.Server.Backmen.Disease;
+
+public sealed class DiseaseInfectionSpread : Job<object>
+{
+    private readonly DiseaseInfectionSpreadEvent _ev;
+    private readonly DiseaseEffectSystem _effectSystem;
+
+    public DiseaseInfectionSpread(DiseaseInfectionSpreadEvent ev, DiseaseEffectSystem effectSystem, double maxTime, CancellationToken cancellation = default) : base(maxTime, cancellation)
+    {
+        _ev = ev;
+        _effectSystem = effectSystem;
+    }
+
+    public DiseaseInfectionSpread(DiseaseInfectionSpreadEvent ev, DiseaseEffectSystem effectSystem, double maxTime, IStopwatch stopwatch, CancellationToken cancellation = default) : base(maxTime, stopwatch, cancellation)
+    {
+        _ev = ev;
+        _effectSystem = effectSystem;
+    }
+
+    protected override async Task<object?> Process()
+    {
+        _effectSystem.DoSpread(_ev.Owner, _ev.Disease, _ev.Range);
+        return null;
+    }
+}