diff --git a/binaries/data/mods/public/simulation/components/UnitAI.js b/binaries/data/mods/public/simulation/components/UnitAI.js
index 300ba321fd..a631e58324 100644
--- a/binaries/data/mods/public/simulation/components/UnitAI.js
+++ b/binaries/data/mods/public/simulation/components/UnitAI.js
@@ -799,12 +799,12 @@ var UnitFsmSpec = {
 			if (cmpTargetUnitAI && cmpTargetUnitAI.IsFormationMember())
 				target = cmpTargetUnitAI.GetFormationController();
 
+			var cmpAttack = Engine.QueryInterface(this.entity, IID_Attack);
 			// Check if we are already in range, otherwise walk there
 			if (!this.CheckTargetAttackRange(target, target))
 			{
 				if (this.TargetIsAlive(target) && this.CheckTargetVisible(target))
 				{
-					var cmpAttack = Engine.QueryInterface(this.entity, IID_Attack);
 					var range = cmpAttack.GetRange(target);
 					this.PushOrderFront("WalkToTargetRange", { "target": target, "min": range.min, "max": range.max }); 
 					return;
@@ -813,7 +813,10 @@ var UnitFsmSpec = {
 				return;
 			}
 			this.CallMemberFunction("Attack", [target, false]); 
-			this.SetNextState("ATTACKING");
+			if (cmpAttack.CanAttackAsFormation())
+				this.SetNextState("ATTACKING");
+			else
+				this.SetNextState("MEMBER");
 		},
 
 		"Order.Garrison": function(msg) {
diff --git a/binaries/data/mods/public/simulation/components/tests/test_UnitAI.js b/binaries/data/mods/public/simulation/components/tests/test_UnitAI.js
index 3947df7e80..c5a5281994 100644
--- a/binaries/data/mods/public/simulation/components/tests/test_UnitAI.js
+++ b/binaries/data/mods/public/simulation/components/tests/test_UnitAI.js
@@ -74,6 +74,7 @@ function TestFormationExiting(mode)
 	AddMock(unit, IID_Position, {
 		GetPosition: function() { return { "x": 0, "y": 0,"z": 0 }; },
 		GetPosition2D: function() { return { "x": 0, "y": 0 }; },
+		GetRotation: function() { return { "y": 0 }; },
 		IsInWorld: function() { return true; },
 	});
 
@@ -124,6 +125,7 @@ function TestFormationExiting(mode)
 		JumpTo: function(x, z) { this.x = x; this.z = z; },
 		GetPosition: function() { return { "x": this.x, "z": this.z }; },
 		GetPosition2D: function() { return { "x": this.x, "y": this.z }; },
+		GetRotation: function() { return { "y": 0 }; },
 		IsInWorld: function() { return true; },
 	});
 
@@ -217,6 +219,8 @@ function TestMoveIntoFormationWhileAttacking()
 	
 		AddMock(unit + i, IID_Position, {
 			GetPosition: function() { return { "x": 0, "z": 0 }; },
+			GetPosition2D: function() { return { "x": 0, "y": 0 }; },
+			GetRotation: function() { return { "y": 0 }; },
 			IsInWorld: function() { return true; },
 		});
 	
@@ -259,6 +263,8 @@ function TestMoveIntoFormationWhileAttacking()
 	AddMock(controller, IID_Position, {
 		JumpTo: function(x, z) { this.x = x; this.z = z; },
 		GetPosition: function() { return { "x": this.x, "z": this.z }; },
+		GetPosition2D: function() { return { "x": this.x, "y": this.z }; },
+		GetRotation: function() { return { "y": 0 }; },
 		IsInWorld: function() { return true; },
 	});
 
@@ -271,6 +277,7 @@ function TestMoveIntoFormationWhileAttacking()
 
 	AddMock(controller, IID_Attack, {
 		GetRange: function() { return {"max":10, "min": 0}; },
+		CanAttackAsFormation: function() { return false },
 	});
 
 	controllerAI.OnCreate();