diff --git a/src/grid.cpp b/src/grid.cpp
index 8bb4923..ce77d55 100644
--- a/src/grid.cpp
+++ b/src/grid.cpp
@@ -86,6 +86,7 @@ namespace envlibcpp {
                 location.removeEntity(entity);
             }
         }
+        entity.setGridId(-1);
     }
 
     bool Grid::isEntityPresent(Entity &entity) {
diff --git a/src/location.cpp b/src/location.cpp
index 25bfae3..7e0f779 100644
--- a/src/location.cpp
+++ b/src/location.cpp
@@ -31,11 +31,17 @@ namespace envlibcpp {
     }
     
     void Location::removeEntity(Entity &entity) {
-        for (auto i = entities.begin(); i != entities.end(); i++) {
-            if (i->getId() == entity.getId()) {
-                entities.erase(i);
+        int index = -1;
+        for (int i = 0; i < getNumEntities(); i++) {
+            if (entities[i].getId() == entity.getId()) {
+                index = i;
+                break;
             }
         }
+        if (index != -1) {
+            entities.erase(entities.begin() + index);
+        }
+        entity.setLocationId(-1);
     }
     
     bool Location::isEntityPresent(Entity &entity) {
diff --git a/src/tests.cpp b/src/tests.cpp
index f25a266..64ffe4d 100644
--- a/src/tests.cpp
+++ b/src/tests.cpp
@@ -11,27 +11,37 @@ using namespace envlibcpp;
 
 void testTemplate() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Template" << std::endl;
-
+    std::cout << "Test 0 - Template" << std::endl;
     assert(true);
     std::cout << "Success" << std::endl;
 }
 
 void testPlacingEntityInLocation() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Placing entity in location" << std::endl;
+    std::cout << "Test 1 - Placing entity in location" << std::endl;
     Entity entity("Daniel");
     Location location(0, 0);
-
     location.addEntity(entity);
-
     assert(location.isEntityPresent(entity) == true);
+    assert(entity.getLocationId() == location.getId());
+    std::cout << "Success" << std::endl;
+}
+
+void testRemovingEntityFromLocation() {
+    std::cout << "---" << std::endl;
+    std::cout << "Test 2 - Removing entity from location" << std::endl;
+    Entity entity("Daniel");
+    Location location(0, 0);
+    location.addEntity(entity);
+    location.removeEntity(entity);
+    assert(!location.isEntityPresent(entity));
+    assert(entity.getLocationId() == -1);
     std::cout << "Success" << std::endl;
 }
 
 void testGeneratingGrid() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Generating grid" << std::endl;
+    std::cout << "Test 3 - Generating grid" << std::endl;
     int size = 4;
     Environment environment("test environment", 4);
     int numLocations = environment.getGrid().getLocations().size();
@@ -42,7 +52,7 @@ void testGeneratingGrid() {
 
 void testRetrievingLocation() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Retrieving location" << std::endl;
+    std::cout << "Test 4 - Retrieving location" << std::endl;
     Grid grid(5);
     Location firstLocation = grid.getFirstLocation();
     assert(firstLocation.getX() == 0 && firstLocation.getY() == 0);
@@ -51,10 +61,11 @@ void testRetrievingLocation() {
 
 void testRetrievingLocationAfterModification() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Retrieving location after modification" << std::endl;
+    std::cout << "Test 5 - Retrieving location after modification" << std::endl;
     Grid grid(5);
     Location firstLocation = grid.getFirstLocation();
     Entity entity = Entity("test");
+    firstLocation = grid.getFirstLocation();
     firstLocation.addEntity(entity);
     assert(firstLocation.getX() == 0 && firstLocation.getY() == 0);
     assert(firstLocation.isEntityPresent(entity));
@@ -63,35 +74,67 @@ void testRetrievingLocationAfterModification() {
 
 void testPlacingEntityInGrid() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Placing entity in grid" << std::endl;
+    std::cout << "Test 6 - Placing entity in grid" << std::endl;
     Entity entity("Daniel");
     Grid grid(4);
-
     grid.addEntity(entity);
-
     assert(grid.isEntityPresent(entity));
+    assert(entity.getGridId() == grid.getId());
+    std::cout << "Success" << std::endl;
+}
+
+void testRemovingEntityFromGrid() {
+    std::cout << "---" << std::endl;
+    std::cout << "Test 7 - Removing entity from grid" << std::endl;
+    Entity entity("Daniel");
+    Grid grid(4);
+    grid.addEntity(entity);
+    grid.removeEntity(entity);
+    assert(!grid.isEntityPresent(entity));
+    assert(entity.getGridId() == -1);
     std::cout << "Success" << std::endl;
 }
 
 void testPlacingEntityInEnvironment() {
     std::cout << "---" << std::endl;
-    std::cout << "Test - Placing entity in environment" << std::endl;
+    std::cout << "Test 8 - Placing entity in environment" << std::endl;
     Entity entity("Daniel");
     Environment environment("Earth", 2);
-
     environment.addEntity(entity);
+    assert(environment.isEntityPresent(entity));
+    assert(entity.getEnvironmentId() == environment.getId());
+    std::cout << "Success" << std::endl;
+}
 
-    assert(environment.isEntityPresent(entity) == true);
+void testRemovingEntityFromEnvironment() {
+    std::cout << "---" << std::endl;
+    std::cout << "Test 9 - Removing entity from environment" << std::endl;
+    Entity entity("Daniel");
+    Environment environment("Earth", 2);
+    environment.addEntity(entity);
+    environment.removeEntity(entity);
+    assert(!environment.isEntityPresent(entity));
+    assert(entity.getEnvironmentId() == -1);
     std::cout << "Success" << std::endl;
 }
 
-int main() {
+void seedRandomNumberGenerator() {
     srand (time (NULL));
+}
+
+int main() {
+    // seed RNG
+    seedRandomNumberGenerator();
+    
+    // tests
     testPlacingEntityInLocation();
+    testRemovingEntityFromLocation();
     testGeneratingGrid();
     testRetrievingLocation();
     testRetrievingLocationAfterModification();
     testPlacingEntityInGrid();
+    testRemovingEntityFromGrid();
     testPlacingEntityInEnvironment();
+    testRemovingEntityFromEnvironment();
     return 0;
 }
\ No newline at end of file
diff --git a/tests b/tests
index 4d6e6e5..75c0684 100644
Binary files a/tests and b/tests differ