Add a test for the weightless status (#2437)

* Add wip test for weightless status

* Expand and fix weightless test

Content.Client/GameObjects/Components/Mobs/ClientStatusEffectsComponent.cs

@@ -35,6 +35,8 @@ namespace Content.Client.GameObjects.Components.Mobs
         [ViewVariables]
         private Dictionary<StatusEffect, CooldownGraphic> _cooldown = new Dictionary<StatusEffect, CooldownGraphic>();
 
+        public override IReadOnlyDictionary<StatusEffect, StatusEffectStatus> Statuses => _status;
+
         /// <summary>
         /// Allows calculating if we need to act due to this component being controlled by the current mob
         /// </summary>

Content.IntegrationTests/Tests/Gravity/WeightlessStatusTests.cs

@@ -0,0 +1,96 @@
+using System.Threading.Tasks;
+using Content.Server.GameObjects.Components.Gravity;
+using Content.Server.GameObjects.EntitySystems;
+using Content.Shared.GameObjects.Components.Gravity;
+using Content.Shared.GameObjects.Components.Mobs;
+using Content.Shared.GameObjects.EntitySystems;
+using Content.Shared.Utility;
+using NUnit.Framework;
+using Robust.Server.Interfaces.Timing;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Interfaces.GameObjects;
+using Robust.Shared.Interfaces.Map;
+using Robust.Shared.Map;
+
+namespace Content.IntegrationTests.Tests.Gravity
+{
+    [TestFixture]
+    [TestOf(typeof(WeightlessSystem))]
+    [TestOf(typeof(GravityGeneratorComponent))]
+    public class WeightlessStatusTests : ContentIntegrationTest
+    {
+        [Test]
+        public async Task WeightlessStatusTest()
+        {
+            var server = StartServer();
+
+            await server.WaitIdleAsync();
+
+            var mapManager = server.ResolveDependency<IMapManager>();
+            var entityManager = server.ResolveDependency<IEntityManager>();
+            var pauseManager = server.ResolveDependency<IPauseManager>();
+            var tileDefinitionManager = server.ResolveDependency<ITileDefinitionManager>();
+
+            IEntity human = null;
+            SharedStatusEffectsComponent statusEffects = null;
+
+            await server.WaitAssertion(() =>
+            {
+                var mapId = mapManager.CreateMap();
+
+                pauseManager.AddUninitializedMap(mapId);
+
+                var gridId = new GridId(1);
+
+                if (!mapManager.TryGetGrid(gridId, out var grid))
+                {
+                    grid = mapManager.CreateGrid(mapId, gridId);
+                }
+
+                var tileDefinition = tileDefinitionManager["underplating"];
+                var tile = new Tile(tileDefinition.TileId);
+                var coordinates = grid.ToCoordinates();
+
+                grid.SetTile(coordinates, tile);
+
+                pauseManager.DoMapInitialize(mapId);
+
+                human = entityManager.SpawnEntity("HumanMob_Content", coordinates);
+
+                Assert.True(human.TryGetComponent(out statusEffects));
+            });
+
+            // Let WeightlessSystem and GravitySystem tick
+            await server.WaitRunTicks(1);
+
+            GravityGeneratorComponent gravityGenerator = null;
+
+            await server.WaitAssertion(() =>
+            {
+                // No gravity without a gravity generator
+                Assert.True(statusEffects.Statuses.ContainsKey(StatusEffect.Weightless));
+
+                gravityGenerator = human.EnsureComponent<GravityGeneratorComponent>();
+            });
+
+            // Let WeightlessSystem and GravitySystem tick
+            await server.WaitRunTicks(1);
+
+            await server.WaitAssertion(() =>
+            {
+                Assert.False(statusEffects.Statuses.ContainsKey(StatusEffect.Weightless));
+
+                // Disable the gravity generator
+                var args = new BreakageEventArgs {Owner = human};
+                gravityGenerator.OnBreak(args);
+            });
+
+            await server.WaitRunTicks(1);
+
+            await server.WaitAssertion(() =>
+            {
+                Assert.False(statusEffects.Statuses.ContainsKey(StatusEffect.Weightless));
+            });
+        }
+    }
+}

Content.Server/GameObjects/Components/Mobs/ServerStatusEffectsComponent.cs

@@ -23,6 +23,8 @@ namespace Content.Server.GameObjects.Components.Mobs
         [ViewVariables]
         private readonly Dictionary<StatusEffect, StatusEffectStatus> _statusEffects = new Dictionary<StatusEffect, StatusEffectStatus>();
 
+        public override IReadOnlyDictionary<StatusEffect, StatusEffectStatus> Statuses => _statusEffects;
+
         protected override void Startup()
         {
             base.Startup();

Content.Server/GameObjects/EntitySystems/WeightlessSystem.cs

@@ -7,6 +7,8 @@ using JetBrains.Annotations;
 using Robust.Shared.GameObjects.Components.Map;
 using Robust.Shared.GameObjects.EntitySystemMessages;
 using Robust.Shared.GameObjects.Systems;
+using Robust.Shared.Interfaces.Map;
+using Robust.Shared.IoC;
 using Robust.Shared.Map;
 using Robust.Shared.Utility;
 
@@ -15,6 +17,8 @@ namespace Content.Server.GameObjects.EntitySystems
     [UsedImplicitly]
     public class WeightlessSystem : EntitySystem, IResettingEntitySystem
     {
+        [Dependency] private readonly IMapManager _mapManager = default!;
+
         private readonly Dictionary<GridId, List<ServerStatusEffectsComponent>> _statuses = new Dictionary<GridId, List<ServerStatusEffectsComponent>>();
 
         public override void Initialize()
@@ -32,10 +36,22 @@ namespace Content.Server.GameObjects.EntitySystems
 
         public void AddStatus(ServerStatusEffectsComponent status)
         {
-            var grid = status.Owner.Transform.GridID;
+            var gridId = status.Owner.Transform.GridID;
-            var statuses = _statuses.GetOrNew(grid);
+            var statuses = _statuses.GetOrNew(gridId);
 
             statuses.Add(status);
+
+            if (_mapManager.TryGetGrid(status.Owner.Transform.GridID, out var grid))
+            {
+                if (grid.HasGravity)
+                {
+                    RemoveWeightless(status);
+                }
+                else
+                {
+                    AddWeightless(status);
+                }
+            }
         }
 
         public void RemoveStatus(ServerStatusEffectsComponent status)

Content.Shared/GameObjects/Components/Mobs/SharedStatusEffectsComponent.cs

@@ -14,6 +14,8 @@ namespace Content.Shared.GameObjects.Components.Mobs
         public override string Name => "StatusEffectsUI";
         public override uint? NetID => ContentNetIDs.STATUSEFFECTS;
 
+        public abstract IReadOnlyDictionary<StatusEffect, StatusEffectStatus> Statuses { get; }
+
         public abstract void ChangeStatusEffectIcon(StatusEffect effect, string icon);
 
         public abstract void ChangeStatusEffect(StatusEffect effect, string icon, ValueTuple<TimeSpan, TimeSpan>? cooldown);