From c45ba3b34860121571eb9a660a76bc82ac0a1a96 Mon Sep 17 00:00:00 2001
From: Chris Nix <chris.nix@gmail.com>
Date: Tue, 22 Dec 2015 22:58:37 +0000
Subject: [PATCH] Process delayed removal and addition of entities without
 exception

---
 .../badlogic/ashley/core/EntityManager.java   |  8 +--
 .../ashley/core/EntityManagerTests.java       | 62 ++++++++++++++++++-
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/ashley/src/com/badlogic/ashley/core/EntityManager.java b/ashley/src/com/badlogic/ashley/core/EntityManager.java
index f622f2c5..c7ddebfd 100644
--- a/ashley/src/com/badlogic/ashley/core/EntityManager.java
+++ b/ashley/src/com/badlogic/ashley/core/EntityManager.java
@@ -23,10 +23,6 @@ public void addEntity(Entity entity){
 	}
 	
 	public void addEntity(Entity entity, boolean delayed){
-		if (entitySet.contains(entity)) {
-			throw new IllegalArgumentException("Entity is already registered " + entity);
-		}
-		
 		if (delayed) {
 			EntityOperation operation = entityOperationPool.obtain();
 			operation.entity = entity;
@@ -113,6 +109,10 @@ protected void removeEntityInternal(Entity entity) {
 	}
 
 	protected void addEntityInternal(Entity entity) {
+		if (entitySet.contains(entity)) {
+			throw new IllegalArgumentException("Entity is already registered " + entity);
+		}
+
 		entities.add(entity);
 		entitySet.add(entity);
 
diff --git a/ashley/tests/com/badlogic/ashley/core/EntityManagerTests.java b/ashley/tests/com/badlogic/ashley/core/EntityManagerTests.java
index cab462ff..22bec101 100644
--- a/ashley/tests/com/badlogic/ashley/core/EntityManagerTests.java
+++ b/ashley/tests/com/badlogic/ashley/core/EntityManagerTests.java
@@ -75,14 +75,34 @@ public void getEntities () {
 	}
 	
 	@Test(expected=IllegalArgumentException.class)
-	public void addEntityTwice () {
+	public void addEntityTwice1 () {
 		 EntityListenerMock listener = new EntityListenerMock();
 		 EntityManager manager = new EntityManager(listener);
 	    Entity entity = new Entity();
 	    manager.addEntity(entity);
 	    manager.addEntity(entity);
 	}
-	
+
+	@Test(expected=IllegalArgumentException.class)
+	public void addEntityTwice2() {
+		EntityListenerMock listener = new EntityListenerMock();
+		EntityManager manager = new EntityManager(listener);
+		Entity entity = new Entity();
+		manager.addEntity(entity, false);
+		manager.addEntity(entity, false);
+	}
+
+	@Test(expected=IllegalArgumentException.class)
+	public void addEntityTwiceDelayed() {
+		EntityListenerMock listener = new EntityListenerMock();
+		EntityManager manager = new EntityManager(listener);
+
+		Entity entity = new Entity();
+		manager.addEntity(entity, true);
+		manager.addEntity(entity, true);
+		manager.processPendingOperations();
+	}
+
 	@Test
 	public void delayedOperationsOrder() {
 		EntityListenerMock listener = new EntityListenerMock();
@@ -108,4 +128,42 @@ public void delayedOperationsOrder() {
 		assertNotEquals(-1, manager.getEntities().indexOf(entityC, true));
 		assertNotEquals(-1, manager.getEntities().indexOf(entityD, true));
 	}
+
+	@Test
+	public void removeAndAddEntityDelayed() {
+		EntityListenerMock listener = new EntityListenerMock();
+		EntityManager manager = new EntityManager(listener);
+
+		Entity entity = new Entity();
+		manager.addEntity(entity, false);     // immediate
+		assertEquals(1, manager.getEntities().size());
+
+		manager.removeEntity(entity, true);   // delayed
+		assertEquals(1, manager.getEntities().size());
+
+		manager.addEntity(entity, true);      // delayed
+		assertEquals(1, manager.getEntities().size());
+
+		manager.processPendingOperations();
+		assertEquals(1, manager.getEntities().size());
+	}
+
+	@Test
+	public void removeAllAndAddEntityDelayed() {
+		EntityListenerMock listener = new EntityListenerMock();
+		EntityManager manager = new EntityManager(listener);
+
+		Entity entity = new Entity();
+		manager.addEntity(entity, false);  // immediate
+		assertEquals(1, manager.getEntities().size());
+
+		manager.removeAllEntities(true);   // delayed
+		assertEquals(1, manager.getEntities().size());
+
+		manager.addEntity(entity, true);   // delayed
+		assertEquals(1, manager.getEntities().size());
+
+		manager.processPendingOperations();
+		assertEquals(1, manager.getEntities().size());
+	}
 }