2 files changed, 17 insertions, 25 deletions
diff --git a/src/server/game/Entities/Creature/Creature.cpp b/src/server/game/Entities/Creature/Creature.cpp
index 00ef3d32a06..650d904db47 100644
--- a/src/server/game/Entities/Creature/Creature.cpp
+++ b/src/server/game/Entities/Creature/Creature.cpp
@@ -478,6 +478,23 @@ void Creature::Update(uint32 diff)
             RemoveUnitMovementFlag(MOVEMENTFLAG_SWIMMING);
     }
 
+    // Set the movement flags if the creature is in that mode. (Only fly if actually in air, only swim if in water, etc)
+    float ground = GetPositionZ();
+    GetMap()->GetWaterOrGroundLevel(GetPositionX(), GetPositionY(), GetPositionZ(), &ground);
+
+    bool isInAir = G3D::fuzzyGt(GetPositionZ(), ground + 0.05f);
+    CreatureTemplate const* cinfo = GetCreatureTemplate();
+
+    if (cinfo->InhabitType & INHABIT_AIR && cinfo->InhabitType & INHABIT_GROUND && isInAir)
+        SetCanFly(true);
+    else if (cinfo->InhabitType & INHABIT_AIR && isInAir)
+        SetDisableGravity(true);
+    else
+    {
+        SetCanFly(false);
+        SetDisableGravity(false);
+    }
+
     switch (m_deathState)
     {
         case JUST_RESPAWNED:
diff --git a/src/server/game/Entities/Unit/Unit.cpp b/src/server/game/Entities/Unit/Unit.cpp
index a9376dca351..62df8c2049c 100644
--- a/src/server/game/Entities/Unit/Unit.cpp
+++ b/src/server/game/Entities/Unit/Unit.cpp
@@ -17375,31 +17375,6 @@ bool Unit::UpdatePosition(float x, float y, float z, float orientation, bool tel
     else if (turn)
         UpdateOrientation(orientation);
 
-    if (Creature* creature = ToCreature())
-    {
-        // Set the movement flags if the creature is in that mode. (Only fly if actually in air, only swim if in water, etc)
-        float ground = z;
-        GetMap()->GetWaterOrGroundLevel(x, y, z, &ground);
-
-        bool isInAir = G3D::fuzzyGt(z, ground + 0.05f);
-        CreatureTemplate const* cinfo = creature->GetCreatureTemplate();
-
-        if (cinfo->InhabitType & INHABIT_AIR && cinfo->InhabitType & INHABIT_GROUND && isInAir)
-            SetCanFly(true);
-        else if (cinfo->InhabitType & INHABIT_AIR && isInAir)
-            SetDisableGravity(true);
-        else
-        {
-            SetCanFly(false);
-            SetDisableGravity(false);
-        }
-
-        if (cinfo->InhabitType & INHABIT_WATER && GetMap()->IsInWater(x, y, z))
-            AddUnitMovementFlag(MOVEMENTFLAG_SWIMMING);
-        else
-            RemoveUnitMovementFlag(MOVEMENTFLAG_SWIMMING);
-    }
-
     // code block for underwater state update
     UpdateUnderwaterState(GetMap(), x, y, z);