Fix mechanism events not being called properly, add test (#2279)

* Add mechanism events when added/removed to/from body/parts

* Change old usages

* Add TODO

* Remove BodyExtensions and IHasBody

* Remove unnecessary extensions and fix wrong event call in mechanism behavior component

* Complete test and fix event calls

Content.Client/GameObjects/Components/Body/Scanner/BodyScannerDisplay.cs

@@ -19,7 +19,7 @@ namespace Content.Client.GameObjects.Components.Body.Scanner
         private IEntity? _currentEntity;
         private IBodyPart? _currentBodyPart;
 
-        private IBody? CurrentBody => _currentEntity?.GetBody();
+        private IBody? CurrentBody => _currentEntity?.GetComponentOrNull<IBody>();
 
         public BodyScannerDisplay(BodyScannerBoundUserInterface owner)
         {

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

@@ -47,7 +47,7 @@ namespace Content.Client.GameObjects.Components.Mobs
                 return;
             }
 
-            if (Owner.TryGetBody(out var body))
+            if (Owner.TryGetComponent(out IBody body))
             {
                 foreach (var part in body.Parts.Values)
                 {

Content.IntegrationTests/Tests/Body/LegTest.cs

@@ -3,10 +3,8 @@ using Content.Server.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Part;
 using Content.Shared.GameObjects.Components.Rotation;
-using Content.Shared.GameObjects.EntitySystems;
 using NUnit.Framework;
 using Robust.Server.GameObjects;
-using Robust.Shared.GameObjects.Systems;
 using Robust.Shared.Interfaces.GameObjects;
 using Robust.Shared.Interfaces.Map;
 using Robust.Shared.IoC;
@@ -36,7 +34,7 @@ namespace Content.IntegrationTests.Tests.Body
                 var entityManager = IoCManager.Resolve<IEntityManager>();
                 var human = entityManager.SpawnEntity("HumanMob_Content", MapCoordinates.Nullspace);
 
-                Assert.That(human.TryGetBody(out var body));
+                Assert.That(human.TryGetComponent(out IBody body));
                 Assert.That(human.TryGetComponent(out appearance));
 
                 Assert.That(!appearance.TryGetData(RotationVisuals.RotationState, out RotationState _));

Content.IntegrationTests/Tests/Body/LungTest.cs

@@ -6,6 +6,7 @@ using Content.Server.GameObjects.Components.Body.Behavior;
 using Content.Server.GameObjects.Components.Body.Circulatory;
 using Content.Server.GameObjects.Components.Metabolism;
 using Content.Shared.Atmos;
+using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using NUnit.Framework;
 using Robust.Server.Interfaces.Maps;
@@ -36,9 +37,10 @@ namespace Content.IntegrationTests.Tests.Body
 
                 var human = entityManager.SpawnEntity("HumanMob_Content", MapCoordinates.Nullspace);
 
-                Assert.True(human.TryGetMechanismBehaviors(out List<LungBehaviorComponent> lungs));
+                Assert.That(human.TryGetComponent(out IBody body));
+                Assert.That(body.TryGetMechanismBehaviors(out List<LungBehaviorComponent> lungs));
                 Assert.That(lungs.Count, Is.EqualTo(1));
-                Assert.True(human.TryGetComponent(out BloodstreamComponent bloodstream));
+                Assert.That(human.TryGetComponent(out BloodstreamComponent bloodstream));
 
                 var gas = new GasMixture(1);
 
@@ -137,7 +139,8 @@ namespace Content.IntegrationTests.Tests.Body
                 var coordinates = new EntityCoordinates(grid.GridEntityId, center);
                 human = entityManager.SpawnEntity("HumanMob_Content", coordinates);
 
-                Assert.True(human.HasMechanismBehavior<LungBehaviorComponent>());
+                Assert.True(human.TryGetComponent(out IBody body));
+                Assert.True(body.HasMechanismBehavior<LungBehaviorComponent>());
                 Assert.True(human.TryGetComponent(out metabolism));
                 Assert.False(metabolism.Suffocating);
             });

Content.IntegrationTests/Tests/Body/MechanismBehaviorEventsTest.cs

@@ -0,0 +1,219 @@
+#nullable enable
+using System.Linq;
+using System.Threading.Tasks;
+using Content.Shared.GameObjects.Components.Body;
+using Content.Shared.GameObjects.Components.Body.Behavior;
+using Content.Shared.GameObjects.Components.Body.Mechanism;
+using Content.Shared.GameObjects.Components.Body.Part;
+using NUnit.Framework;
+using Robust.Shared.GameObjects;
+using Robust.Shared.Interfaces.GameObjects;
+using Robust.Shared.Interfaces.Map;
+using Robust.Shared.IoC;
+using Robust.Shared.Map;
+
+namespace Content.IntegrationTests.Tests.Body
+{
+    [TestFixture]
+    [TestOf(typeof(SharedBodyComponent))]
+    [TestOf(typeof(SharedBodyPartComponent))]
+    [TestOf(typeof(SharedMechanismComponent))]
+    [TestOf(typeof(MechanismBehaviorComponent))]
+    public class MechanismBehaviorEventsTest : ContentIntegrationTest
+    {
+        [RegisterComponent]
+        private class TestBehaviorComponent : MechanismBehaviorComponent
+        {
+            public override string Name => nameof(MechanismBehaviorEventsTest) + "TestBehavior";
+
+            public bool WasAddedToBody;
+            public bool WasAddedToPart;
+            public bool WasAddedToPartInBody;
+            public bool WasRemovedFromBody;
+            public bool WasRemovedFromPart;
+            public bool WasRemovedFromPartInBody;
+
+            public override void Update(float frameTime) { }
+
+            public bool AllAdded()
+            {
+                return WasAddedToBody && WasAddedToPart && WasAddedToPartInBody;
+            }
+
+            public bool AllRemoved()
+            {
+                return WasRemovedFromBody && WasRemovedFromPart && WasRemovedFromPartInBody;
+            }
+
+            public bool NoAdded()
+            {
+                return !WasAddedToBody && !WasAddedToPart && !WasAddedToPartInBody;
+            }
+
+            public bool NoRemoved()
+            {
+                return !WasRemovedFromBody && !WasRemovedFromPart && !WasRemovedFromPartInBody;
+            }
+
+            public void ResetAdded()
+            {
+                WasAddedToBody = false;
+                WasAddedToPart = false;
+                WasAddedToPartInBody = false;
+            }
+
+            public void ResetRemoved()
+            {
+                WasRemovedFromBody = false;
+                WasRemovedFromPart = false;
+                WasRemovedFromPartInBody = false;
+            }
+
+            public void ResetAll()
+            {
+                ResetAdded();
+                ResetRemoved();
+            }
+
+            protected override void OnAddedToBody()
+            {
+                base.OnAddedToBody();
+
+                WasAddedToBody = true;
+            }
+
+            protected override void OnAddedToPart()
+            {
+                base.OnAddedToPart();
+
+                WasAddedToPart = true;
+            }
+
+            protected override void OnAddedToPartInBody()
+            {
+                base.OnAddedToPartInBody();
+
+                WasAddedToPartInBody = true;
+            }
+
+            protected override void OnRemovedFromBody(IBody old)
+            {
+                base.OnRemovedFromBody(old);
+
+                WasRemovedFromBody = true;
+            }
+
+            protected override void OnRemovedFromPart(IBodyPart old)
+            {
+                base.OnRemovedFromPart(old);
+
+                WasRemovedFromPart = true;
+            }
+
+            protected override void OnRemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart)
+            {
+                base.OnRemovedFromPartInBody(oldBody, oldPart);
+
+                WasRemovedFromPartInBody = true;
+            }
+        }
+
+        [Test]
+        public async Task EventsTest()
+        {
+            var server = StartServerDummyTicker(new ServerContentIntegrationOption
+            {
+                ContentBeforeIoC = () =>
+                {
+                    IoCManager.Resolve<IComponentFactory>().Register<TestBehaviorComponent>();
+                }
+            });
+
+            await server.WaitAssertion(() =>
+            {
+                var mapManager = IoCManager.Resolve<IMapManager>();
+
+                var mapId = new MapId(0);
+                mapManager.CreateNewMapEntity(mapId);
+
+                var entityManager = IoCManager.Resolve<IEntityManager>();
+                var human = entityManager.SpawnEntity("HumanMob_Content", MapCoordinates.Nullspace);
+
+                Assert.That(human.TryGetComponent(out IBody? body));
+                Assert.NotNull(body);
+
+                var centerPart = body!.CenterPart();
+                Assert.NotNull(centerPart);
+
+                Assert.That(body.TryGetSlot(centerPart!, out var centerSlot));
+                Assert.NotNull(centerSlot);
+
+                var mechanism = centerPart!.Mechanisms.First();
+                Assert.NotNull(mechanism);
+
+                var component = mechanism.Owner.AddComponent<TestBehaviorComponent>();
+                Assert.False(component.WasAddedToBody);
+                Assert.False(component.WasAddedToPart);
+                Assert.That(component.WasAddedToPartInBody);
+                Assert.That(component.NoRemoved);
+
+                component.ResetAll();
+
+                Assert.That(component.NoAdded);
+                Assert.That(component.NoRemoved);
+
+                centerPart.RemoveMechanism(mechanism);
+
+                Assert.That(component.NoAdded);
+                Assert.False(component.WasRemovedFromBody);
+                Assert.False(component.WasRemovedFromPart);
+                Assert.That(component.WasRemovedFromPartInBody);
+
+                component.ResetAll();
+
+                centerPart.TryAddMechanism(mechanism, true);
+
+                Assert.False(component.WasAddedToBody);
+                Assert.False(component.WasAddedToPart);
+                Assert.That(component.WasAddedToPartInBody);
+                Assert.That(component.NoRemoved());
+
+                component.ResetAll();
+
+                body.RemovePart(centerPart, true);
+
+                Assert.That(component.NoAdded);
+                Assert.That(component.WasRemovedFromBody);
+                Assert.False(component.WasRemovedFromPart);
+                Assert.False(component.WasRemovedFromPartInBody);
+
+                component.ResetAll();
+
+                centerPart.RemoveMechanism(mechanism);
+
+                Assert.That(component.NoAdded);
+                Assert.False(component.WasRemovedFromBody);
+                Assert.That(component.WasRemovedFromPart);
+                Assert.False(component.WasRemovedFromPartInBody);
+
+                component.ResetAll();
+
+                centerPart.TryAddMechanism(mechanism, true);
+
+                Assert.False(component.WasAddedToBody);
+                Assert.That(component.WasAddedToPart);
+                Assert.False(component.WasAddedToPartInBody);
+                Assert.That(component.NoRemoved);
+
+                component.ResetAll();
+
+                body.TryAddPart(centerSlot!, centerPart, true);
+
+                Assert.That(component.WasAddedToBody);
+                Assert.False(component.WasAddedToPart);
+                Assert.False(component.WasAddedToPartInBody);
+                Assert.That(component.NoRemoved);
+            });
+        }
+    }
+}

Content.IntegrationTests/Tests/BuckleTest.cs

@@ -3,11 +3,9 @@ using Content.Server.GameObjects.Components.Buckle;
 using Content.Server.GameObjects.Components.GUI;
 using Content.Server.GameObjects.Components.Items.Storage;
 using Content.Server.GameObjects.Components.Strap;
-using Content.Shared.Damage;
 using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Part;
 using Content.Shared.GameObjects.Components.Buckle;
-using Content.Shared.GameObjects.Components.Damage;
 using Content.Shared.GameObjects.EntitySystems;
 using Content.Shared.Utility;
 using NUnit.Framework;
@@ -210,7 +208,7 @@ namespace Content.IntegrationTests.Tests
                 Assert.True(human.TryGetComponent(out buckle));
                 Assert.True(chair.TryGetComponent(out strap));
                 Assert.True(human.TryGetComponent(out hands));
-                Assert.True(human.TryGetBody(out body));
+                Assert.True(human.TryGetComponent(out body));
 
                 // Buckle
                 Assert.True(buckle.TryBuckle(human, chair));

Content.IntegrationTests/Tests/GameObjects/Components/ActionBlocking/CuffUnitTest.cs

@@ -53,7 +53,7 @@ namespace Content.IntegrationTests.Tests.GameObjects.Components.ActionBlocking
                 // Test for components existing
                 Assert.True(human.TryGetComponent(out cuffed!), $"Human has no {nameof(CuffableComponent)}");
                 Assert.True(human.TryGetComponent(out hands!), $"Human has no {nameof(HandsComponent)}");
-                Assert.True(human.TryGetBody(out body!), $"Human has no {nameof(IBody)}");
+                Assert.True(human.TryGetComponent(out body!), $"Human has no {nameof(IBody)}");
                 Assert.True(cuffs.TryGetComponent(out handcuff!), $"Handcuff has no {nameof(HandcuffComponent)}");
                 Assert.True(cables.TryGetComponent(out cableHandcuff!), $"Cablecuff has no {nameof(HandcuffComponent)}");
 

Content.Server/GameObjects/Components/Body/BodyCommands.cs

@@ -169,7 +169,7 @@ namespace Content.Server.GameObjects.Components.Body
                 return;
             }
 
-            if (!player.AttachedEntity.TryGetBody(out var body))
+            if (!player.AttachedEntity.TryGetComponent(out IBody? body))
             {
                 var random = IoCManager.Resolve<IRobustRandom>();
                 var text = $"You have no body{(random.Prob(0.2f) ? " and you must scream." : ".")}";
@@ -216,7 +216,7 @@ namespace Content.Server.GameObjects.Components.Body
                 return;
             }
 
-            if (!player.AttachedEntity.TryGetBody(out var body))
+            if (!player.AttachedEntity.TryGetComponent(out IBody? body))
             {
                 var random = IoCManager.Resolve<IRobustRandom>();
                 var text = $"You have no body{(random.Prob(0.2f) ? " and you must scream." : ".")}";

Content.Server/GameObjects/Components/Body/MechanismComponent.cs

@@ -47,7 +47,7 @@ namespace Content.Server.GameObjects.Components.Body
             PerformerCache = null;
             BodyCache = null;
 
-            if (eventArgs.Target.TryGetBody(out var body))
+            if (eventArgs.Target.TryGetComponent(out IBody? body))
             {
                 SendBodyPartListToUser(eventArgs, body);
             }

Content.Server/GameObjects/Components/Body/Part/BodyPartComponent.cs

@@ -104,7 +104,7 @@ namespace Content.Server.GameObjects.Components.Body.Part
             _surgeonCache = null;
             _owningBodyCache = null;
 
-            if (eventArgs.Target.TryGetBody(out var body))
+            if (eventArgs.Target.TryGetComponent(out IBody? body))
             {
                 SendSlots(eventArgs, body);
             }

Content.Server/GameObjects/Components/Body/SurgeryToolComponent.cs

@@ -68,7 +68,7 @@ namespace Content.Server.GameObjects.Components.Body
             _callbackCache = null;
 
             // Attempt surgery on a body by sending a list of operable parts for the client to choose from
-            if (eventArgs.Target.TryGetBody(out var body))
+            if (eventArgs.Target.TryGetComponent(out IBody? body))
             {
                 // Create dictionary to send to client (text to be shown : data sent back if selected)
                 var toSend = new Dictionary<string, int>();

Content.Server/GameObjects/Components/Chemistry/PillComponent.cs

@@ -3,6 +3,7 @@ using Content.Server.GameObjects.Components.Body.Behavior;
 using Content.Server.GameObjects.Components.Nutrition;
 using Content.Server.GameObjects.Components.Utensil;
 using Content.Shared.Chemistry;
+using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using Content.Shared.Interfaces;
 using Content.Shared.Interfaces.GameObjects.Components;
@@ -82,7 +83,8 @@ namespace Content.Server.GameObjects.Components.Chemistry
 
             var trueTarget = target ?? user;
 
-            if (!trueTarget.TryGetMechanismBehaviors<StomachBehaviorComponent>(out var stomachs))
+            if (!trueTarget.TryGetComponent(out IBody body) ||
+                !body.TryGetMechanismBehaviors<StomachBehaviorComponent>(out var stomachs))
             {
                 return false;
             }

Content.Server/GameObjects/Components/Metabolism/MetabolismComponent.cs

@@ -1,4 +1,5 @@
-using System;
+#nullable enable
+using System;
 using System.Collections.Generic;
 using System.Linq;
 using Content.Server.Atmos;
@@ -8,12 +9,14 @@ using Content.Server.GameObjects.Components.Temperature;
 using Content.Shared.Atmos;
 using Content.Shared.Chemistry;
 using Content.Shared.Damage;
+using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using Content.Shared.GameObjects.Components.Damage;
 using Content.Shared.GameObjects.EntitySystems;
 using Content.Shared.Interfaces;
 using Content.Shared.Interfaces.Chemistry;
 using Robust.Shared.GameObjects;
+using Robust.Shared.GameObjects.ComponentDependencies;
 using Robust.Shared.Interfaces.GameObjects;
 using Robust.Shared.IoC;
 using Robust.Shared.Localization;
@@ -29,6 +32,8 @@ namespace Content.Server.GameObjects.Components.Metabolism
         [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
         [Dependency] private readonly IEntityManager _entityManager = default!;
 
+        [ComponentDependency] private readonly IBody? _body = default!;
+
         public override string Name => "Metabolism";
 
         private float _accumulatedFrameTime;
@@ -38,11 +43,11 @@ namespace Content.Server.GameObjects.Components.Metabolism
 
         [ViewVariables(VVAccess.ReadWrite)] private int _suffocationDamage;
 
-        [ViewVariables] public Dictionary<Gas, float> NeedsGases { get; set; }
+        [ViewVariables] public Dictionary<Gas, float> NeedsGases { get; set; } = new Dictionary<Gas, float>();
 
-        [ViewVariables] public Dictionary<Gas, float> ProducesGases { get; set; }
+        [ViewVariables] public Dictionary<Gas, float> ProducesGases { get; set; } = new Dictionary<Gas, float>();
 
-        [ViewVariables] public Dictionary<Gas, float> DeficitGases { get; set; }
+        [ViewVariables] public Dictionary<Gas, float> DeficitGases { get; set; } = new Dictionary<Gas, float>();
 
         /// <summary>
         /// Heat generated due to metabolism. It's generated via metabolism
@@ -176,11 +181,18 @@ namespace Content.Server.GameObjects.Components.Metabolism
 
         private void ProcessGases(float frameTime)
         {
-            if (!Owner.TryGetComponent(out BloodstreamComponent bloodstream))
+            if (!Owner.TryGetComponent(out BloodstreamComponent? bloodstream))
             {
                 return;
             }
 
+            if (_body == null)
+            {
+                return;
+            }
+
+            var lungs = _body.GetMechanismBehaviors<LungBehaviorComponent>().ToArray();
+
             var needs = NeedsAndDeficit(frameTime);
             var used = 0f;
             foreach (var (gas, amountNeeded) in needs)
@@ -191,16 +203,13 @@ namespace Content.Server.GameObjects.Components.Metabolism
                 if (bloodstreamAmount < amountNeeded)
                 {
                     // Panic inhale
-                    if (Owner.TryGetMechanismBehaviors(out List<LungBehaviorComponent> lungs))
+                    foreach (var lung in lungs)
                     {
-                        foreach (var lung in lungs)
+                        lung.Gasp();
-                        {
-                            lung.Gasp();
-                        }
-
-                        bloodstreamAmount = bloodstream.Air.GetMoles(gas);
                     }
 
+                    bloodstreamAmount = bloodstream.Air.GetMoles(gas);
+
                     deficit = Math.Max(0, amountNeeded - bloodstreamAmount);
 
                     if (deficit > 0)
@@ -238,7 +247,7 @@ namespace Content.Server.GameObjects.Components.Metabolism
         /// <param name="frameTime"></param>
         private void ProcessThermalRegulation(float frameTime)
         {
-            if (!Owner.TryGetComponent(out TemperatureComponent temperatureComponent)) return;
+            if (!Owner.TryGetComponent(out TemperatureComponent? temperatureComponent)) return;
             temperatureComponent.ReceiveHeat(MetabolismHeat);
             temperatureComponent.RemoveHeat(RadiatedHeat);
 
@@ -308,7 +317,7 @@ namespace Content.Server.GameObjects.Components.Metabolism
         /// <param name="frameTime">The time since the last metabolism tick in seconds.</param>
         private void ProcessNutrients(float frameTime)
         {
-            if (!Owner.TryGetComponent(out BloodstreamComponent bloodstream))
+            if (!Owner.TryGetComponent(out BloodstreamComponent? bloodstream))
             {
                 return;
             }
@@ -377,7 +386,7 @@ namespace Content.Server.GameObjects.Components.Metabolism
         {
             Suffocating = true;
 
-            if (!Owner.TryGetComponent(out IDamageableComponent damageable))
+            if (!Owner.TryGetComponent(out IDamageableComponent? damageable))
             {
                 return;
             }

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

@@ -16,7 +16,7 @@ namespace Content.Server.GameObjects.Components.Mobs
             {
                 base.Appearance = value;
 
-                if (Owner.TryGetBody(out var body))
+                if (Owner.TryGetComponent(out IBody body))
                 {
                     foreach (var part in body.Parts.Values)
                     {
@@ -35,7 +35,7 @@ namespace Content.Server.GameObjects.Components.Mobs
         {
             base.Startup();
 
-            if (Appearance != null && Owner.TryGetBody(out var body))
+            if (Appearance != null && Owner.TryGetComponent(out IBody body))
             {
                 foreach (var part in body.Parts.Values)
                 {

Content.Server/GameObjects/Components/Nutrition/DrinkComponent.cs

@@ -5,6 +5,7 @@ using Content.Server.GameObjects.Components.Fluids;
 using Content.Server.GameObjects.EntitySystems;
 using Content.Shared.Audio;
 using Content.Shared.Chemistry;
+using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using Content.Shared.GameObjects.Components.Nutrition;
 using Content.Shared.GameObjects.EntitySystems;
@@ -151,7 +152,8 @@ namespace Content.Server.GameObjects.Components.Nutrition
                 return false;
             }
 
-            if (!target.TryGetMechanismBehaviors<StomachBehaviorComponent>(out var stomachs))
+            if (!target.TryGetComponent(out IBody body) ||
+                !body.TryGetMechanismBehaviors<StomachBehaviorComponent>(out var stomachs))
             {
                 return false;
             }

Content.Server/GameObjects/Components/Nutrition/FoodComponent.cs

@@ -7,6 +7,7 @@ using Content.Server.GameObjects.Components.GUI;
 using Content.Server.GameObjects.Components.Items.Storage;
 using Content.Server.GameObjects.Components.Utensil;
 using Content.Shared.Chemistry;
+using Content.Shared.GameObjects.Components.Body;
 using Content.Shared.GameObjects.Components.Body.Behavior;
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using Content.Shared.GameObjects.Components.Utensil;
@@ -131,7 +132,8 @@ namespace Content.Server.GameObjects.Components.Nutrition
 
             var trueTarget = target ?? user;
 
-            if (!trueTarget.TryGetMechanismBehaviors<SharedStomachBehaviorComponent>(out var stomachs))
+            if (!trueTarget.TryGetComponent(out IBody? body) ||
+                !body.TryGetMechanismBehaviors<SharedStomachBehaviorComponent>(out var stomachs))
             {
                 return false;
             }

Content.Shared/GameObjects/Components/Body/Behavior/IMechanismBehavior.cs

@@ -1,11 +1,14 @@
 #nullable enable
 using Content.Shared.GameObjects.Components.Body.Mechanism;
 using Content.Shared.GameObjects.Components.Body.Part;
+using Robust.Shared.Interfaces.GameObjects;
 
 namespace Content.Shared.GameObjects.Components.Body.Behavior
 {
-    public interface IMechanismBehavior : IHasBody
+    public interface IMechanismBehavior : IComponent
     {
+        IBody? Body { get; }
+
         IBodyPart? Part { get; }
 
         /// <summary>
@@ -33,6 +36,14 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
         /// </summary>
         void AddedToPart();
 
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is added to a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, adding a brain to a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToPartInBody();
+
         /// <summary>
         ///     Called when the parent <see cref="IBodyPart"/> is removed from a
         ///     <see cref="IBody"/>.
@@ -50,14 +61,6 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
         /// </summary>
         void RemovedFromPart(IBodyPart old);
 
-        /// <summary>
-        ///     Called when the parent <see cref="IMechanism"/> is added to a
-        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
-        ///     For instance, adding a brain to a head that is attached to a body.
-        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
-        /// </summary>
-        void AddedToPartInBody();
-
         /// <summary>
         ///     Called when the parent <see cref="IMechanism"/> is removed from a
         ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.

Content.Shared/GameObjects/Components/Body/Behavior/MechanismBehaviorComponent.cs

@@ -13,6 +13,25 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
 
         public IMechanism? Mechanism => Owner.GetComponentOrNull<IMechanism>();
 
+        protected override void Startup()
+        {
+            base.Startup();
+
+            if (Part == null)
+            {
+                return;
+            }
+
+            if (Body == null)
+            {
+                AddedToPart();
+            }
+            else
+            {
+                AddedToPartInBody();
+            }
+        }
+
         public abstract void Update(float frameTime);
 
         public void AddedToBody()
@@ -37,7 +56,7 @@ namespace Content.Shared.GameObjects.Components.Body.Behavior
 
         public void AddedToPartInBody()
         {
-            OnAddedToPart();
+            OnAddedToPartInBody();
         }
 
         public void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart)

Content.Shared/GameObjects/Components/Body/BodyExtensions.cs

@@ -1,29 +0,0 @@
-#nullable enable
-using System.Diagnostics.CodeAnalysis;
-using Robust.Shared.Interfaces.GameObjects;
-
-namespace Content.Shared.GameObjects.Components.Body
-{
-    public static class BodyExtensions
-    {
-        public static T? GetBody<T>(this IEntity entity) where T : class, IBody
-        {
-            return entity.GetComponentOrNull<T>();
-        }
-
-        public static bool TryGetBody<T>(this IEntity entity, [NotNullWhen(true)] out T? body) where T : class, IBody
-        {
-            return (body = entity.GetBody<T>()) != null;
-        }
-
-        public static IBody? GetBody(this IEntity entity)
-        {
-            return entity.GetComponentOrNull<IBody>();
-        }
-
-        public static bool TryGetBody(this IEntity entity, [NotNullWhen(true)] out IBody? body)
-        {
-            return (body = entity.GetBody()) != null;
-        }
-    }
-}

Content.Shared/GameObjects/Components/Body/IHasBody.cs

@@ -1,13 +0,0 @@
-#nullable enable
-using Robust.Shared.Interfaces.GameObjects;
-
-namespace Content.Shared.GameObjects.Components.Body
-{
-    public interface IHasBody : IComponent
-    {
-        /// <summary>
-        ///     The body that this component is currently a part of, if any.
-        /// </summary>
-        IBody? Body { get; }
-    }
-}

Content.Shared/GameObjects/Components/Body/Mechanism/IMechanism.cs

@@ -1,10 +1,13 @@
 #nullable enable
 using Content.Shared.GameObjects.Components.Body.Part;
+using Robust.Shared.Interfaces.GameObjects;
 
 namespace Content.Shared.GameObjects.Components.Body.Mechanism
 {
-    public interface IMechanism : IHasBody
+    public interface IMechanism : IComponent
     {
+        IBody? Body { get; }
+
         IBodyPart? Part { get; set; }
 
         /// <summary>
@@ -70,6 +73,14 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
         /// </summary>
         void AddedToPart();
 
+        /// <summary>
+        ///     Called when the parent <see cref="IMechanism"/> is added to a
+        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
+        ///     For instance, adding a brain to a head that is attached to a body.
+        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
+        /// </summary>
+        void AddedToPartInBody();
+
         /// <summary>
         ///     Called when the parent <see cref="IBodyPart"/> is removed from a
         ///     <see cref="IBody"/>.
@@ -87,14 +98,6 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
         /// </summary>
         void RemovedFromPart(IBodyPart old);
 
-        /// <summary>
-        ///     Called when the parent <see cref="IMechanism"/> is added to a
-        ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.
-        ///     For instance, adding a brain to a head that is attached to a body.
-        ///     DO NOT CALL THIS DIRECTLY FROM OUTSIDE BODY SYSTEM CODE!
-        /// </summary>
-        void AddedToPartInBody();
-
         /// <summary>
         ///     Called when the parent <see cref="IMechanism"/> is removed from a
         ///     <see cref="IBodyPart"/> that is attached to a <see cref="IBody"/>.

Content.Shared/GameObjects/Components/Body/Mechanism/MechanismExtensions.cs

@@ -3,26 +3,29 @@ using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Content.Shared.GameObjects.Components.Body.Behavior;
-using Robust.Shared.Interfaces.GameObjects;
+using Content.Shared.GameObjects.Components.Body.Part;
 
 namespace Content.Shared.GameObjects.Components.Body.Mechanism
 {
     public static class MechanismExtensions
     {
-        public static bool HasMechanismBehavior<T>(this IEntity entity) where T : IMechanismBehavior
+        public static bool HasMechanismBehavior<T>(this IBody body)
         {
-            // TODO BODY optimize
+            return body.Parts.Values.Any(p => p.HasMechanismBehavior<T>());
-            return entity.TryGetBody(out var body) &&
-                   body.Parts.Values.Any(p => p.Mechanisms.Any(m => m.Owner.HasComponent<T>()));
         }
 
-        public static IEnumerable<IMechanismBehavior> GetMechanismBehaviors(this IEntity entity)
+        public static bool HasMechanismBehavior<T>(this IBodyPart part)
         {
-            if (!entity.TryGetBody(out var body))
+            return part.Mechanisms.Any(m => m.Owner.HasComponent<T>());
-            {
+        }
-                yield break;
-            }
 
+        public static bool HasMechanismBehavior<T>(this IMechanism mechanism)
+        {
+            return mechanism.Owner.HasComponent<T>();
+        }
+
+        public static IEnumerable<IMechanismBehavior> GetMechanismBehaviors(this IBody body)
+        {
             foreach (var part in body.Parts.Values)
             foreach (var mechanism in part.Mechanisms)
             foreach (var behavior in mechanism.Owner.GetAllComponents<IMechanismBehavior>())
@@ -31,10 +34,10 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             }
         }
 
-        public static bool TryGetMechanismBehaviors(this IEntity entity,
+        public static bool TryGetMechanismBehaviors(this IBody body,
             [NotNullWhen(true)] out List<IMechanismBehavior>? behaviors)
         {
-            behaviors = entity.GetMechanismBehaviors().ToList();
+            behaviors = body.GetMechanismBehaviors().ToList();
 
             if (behaviors.Count == 0)
             {
@@ -45,13 +48,8 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             return true;
         }
 
-        public static IEnumerable<T> GetMechanismBehaviors<T>(this IEntity entity) where T : class, IMechanismBehavior
+        public static IEnumerable<T> GetMechanismBehaviors<T>(this IBody body) where T : class, IMechanismBehavior
         {
-            if (!entity.TryGetBody(out var body))
-            {
-                yield break;
-            }
-
             foreach (var part in body.Parts.Values)
             foreach (var mechanism in part.Mechanisms)
             {
@@ -62,7 +60,7 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             }
         }
 
-        public static bool TryGetMechanismBehaviors<T>(this IEntity entity, [NotNullWhen(true)] out List<T>? behaviors)
+        public static bool TryGetMechanismBehaviors<T>(this IBody entity, [NotNullWhen(true)] out List<T>? behaviors)
             where T : class, IMechanismBehavior
         {
             behaviors = entity.GetMechanismBehaviors<T>().ToList();

Content.Shared/GameObjects/Components/Body/Mechanism/SharedMechanismComponent.cs

@@ -1,5 +1,6 @@
 #nullable enable
 using System.Collections.Generic;
+using Content.Shared.GameObjects.Components.Body.Behavior;
 using Content.Shared.GameObjects.Components.Body.Part;
 using Robust.Shared.GameObjects;
 using Robust.Shared.Interfaces.GameObjects;
@@ -39,12 +40,26 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
 
                 if (old != null)
                 {
-                    OnRemovedFromPart(old);
+                    if (old.Body == null)
+                    {
+                        RemovedFromPart(old);
+                    }
+                    else
+                    {
+                        RemovedFromPartInBody(old.Body, old);
+                    }
                 }
 
                 if (value != null)
                 {
-                    OnAddedToPart();
+                    if (value.Body == null)
+                    {
+                        AddedToPart();
+                    }
+                    else
+                    {
+                        AddedToPartInBody();
+                    }
                 }
             }
         }
@@ -94,7 +109,7 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
 
             OnAddedToBody();
 
-            foreach (var behavior in Owner.GetMechanismBehaviors())
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
             {
                 behavior.AddedToBody();
             }
@@ -104,7 +119,7 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
         {
             OnRemovedFromBody(old);
 
-            foreach (var behavior in Owner.GetMechanismBehaviors())
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
             {
                 behavior.RemovedFromBody(old);
             }
@@ -117,23 +132,12 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             Owner.Transform.AttachParent(Part!.Owner);
             OnAddedToPart();
 
-            foreach (var behavior in Owner.GetMechanismBehaviors())
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
             {
                 behavior.AddedToPart();
             }
         }
 
-        public void RemovedFromPart(IBodyPart old)
-        {
-            Owner.Transform.AttachToGridOrMap();
-            OnRemovedFromPart(old);
-
-            foreach (var behavior in Owner.GetMechanismBehaviors())
-            {
-                behavior.RemovedFromPart(old);
-            }
-        }
-
         public void AddedToPartInBody()
         {
             DebugTools.AssertNotNull(Body);
@@ -142,18 +146,29 @@ namespace Content.Shared.GameObjects.Components.Body.Mechanism
             Owner.Transform.AttachParent(Part!.Owner);
             OnAddedToPartInBody();
 
-            foreach (var behavior in Owner.GetMechanismBehaviors())
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
             {
                 behavior.AddedToPartInBody();
             }
         }
 
+        public void RemovedFromPart(IBodyPart old)
+        {
+            Owner.Transform.AttachToGridOrMap();
+            OnRemovedFromPart(old);
+
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
+            {
+                behavior.RemovedFromPart(old);
+            }
+        }
+
         public void RemovedFromPartInBody(IBody? oldBody, IBodyPart? oldPart)
         {
             Owner.Transform.AttachToGridOrMap();
             OnRemovedFromPartInBody();
 
-            foreach (var behavior in Owner.GetMechanismBehaviors())
+            foreach (var behavior in Owner.GetAllComponents<IMechanismBehavior>())
             {
                 behavior.RemovedFromPartInBody(oldBody, oldPart);
             }

Content.Shared/GameObjects/Components/Body/Part/IBodyPart.cs

@@ -7,7 +7,7 @@ using Robust.Shared.Map;
 
 namespace Content.Shared.GameObjects.Components.Body.Part
 {
-    public interface IBodyPart : IHasBody, IBodyPartContainer
+    public interface IBodyPart : IComponent, IBodyPartContainer
     {
         new IBody? Body { get; set; }
 

Content.Shared/GameObjects/Components/Body/Part/SharedBodyPartComponent.cs

@@ -114,15 +114,6 @@ namespace Content.Shared.GameObjects.Components.Body.Part
             mechanism.Part = this;
             SizeUsed += mechanism.Size;
 
-            if (Body == null)
-            {
-                mechanism.AddedToPart();
-            }
-            else
-            {
-                mechanism.AddedToPartInBody();
-            }
-
             Dirty();
         }
 
@@ -132,15 +123,6 @@ namespace Content.Shared.GameObjects.Components.Body.Part
             mechanism.Part = null;
             SizeUsed -= mechanism.Size;
 
-            if (Body == null)
-            {
-                mechanism.RemovedFromPart(this);
-            }
-            else
-            {
-                mechanism.RemovedFromPartInBody(Body, this);
-            }
-
             Dirty();
         }