diff --git a/lib/spells/effects/Heal.cpp b/lib/spells/effects/Heal.cpp
index 75b4139a6..20db55515 100644
--- a/lib/spells/effects/Heal.cpp
+++ b/lib/spells/effects/Heal.cpp
@@ -16,7 +16,7 @@
 #include "../../NetPacks.h"
 #include "../../battle/IBattleState.h"
 #include "../../battle/CBattleInfoCallback.h"
-#include "../../battle/CUnitState.h"
+#include "../../battle/Unit.h"
 #include "../../serializer/JsonSerializeFormat.h"
 
 
@@ -120,7 +120,7 @@ void Heal::prepareHealEffect(int64_t value, BattleUnitsChanged & pack, RNG & rng
 		{
 			auto unitHPgained = m->applySpellBonus(value, unit);
 
-			auto state = unit->acquireState();
+			auto state = unit->acquire();
 			state->heal(unitHPgained, healLevel, healPower);
 
 			if(unitHPgained > 0)
diff --git a/test/spells/effects/HealTest.cpp b/test/spells/effects/HealTest.cpp
index 00dd8849b..9e03f4681 100644
--- a/test/spells/effects/HealTest.cpp
+++ b/test/spells/effects/HealTest.cpp
@@ -286,7 +286,7 @@ TEST_P(HealApplyTest, Heals)
 	GTEST_ASSERT_EQ(targetUnitState->getAvailableHealth(), unitAmount * unitHP / 2 + 1);
 	GTEST_ASSERT_EQ(targetUnitState->getFirstHPleft(), 1);
 
-	EXPECT_CALL(targetUnit, acquireState()).WillOnce(Return(targetUnitState));
+	EXPECT_CALL(targetUnit, acquire()).WillOnce(Return(targetUnitState));
 
 	EXPECT_CALL(*battleFake, setUnitState(Eq(unitId), _, Gt(0))).Times(1);
 
diff --git a/test/spells/effects/SacrificeTest.cpp b/test/spells/effects/SacrificeTest.cpp
index c9adc966b..b95d6e84a 100644
--- a/test/spells/effects/SacrificeTest.cpp
+++ b/test/spells/effects/SacrificeTest.cpp
@@ -198,7 +198,7 @@ TEST_F(SacrificeApplyTest, ResurrectsTarget)
 		targetUnitState->damage(initialDmg);
 	}
 
-	EXPECT_CALL(targetUnit, acquireState()).WillOnce(Return(targetUnitState));
+	EXPECT_CALL(targetUnit, acquire()).WillOnce(Return(targetUnitState));
 
 	setupDefaultRNG();