From f453b82a582efd2a9328f1d005a6d3af772707cb Mon Sep 17 00:00:00 2001
From: Yohann D'ANELLO <yohann.danello@gmail.com>
Date: Fri, 11 Dec 2020 18:33:47 +0100
Subject: [PATCH] Test clicking on the screen

---
 squirrelbattle/entities/items.py  |  4 ++--
 squirrelbattle/interfaces.py      |  3 ++-
 squirrelbattle/tests/game_test.py | 27 +++++++++++++++++++++++++++
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/squirrelbattle/entities/items.py b/squirrelbattle/entities/items.py
index e90ec32..a1f3bd4 100644
--- a/squirrelbattle/entities/items.py
+++ b/squirrelbattle/entities/items.py
@@ -28,7 +28,7 @@ class Item(Entity):
         """
         if self.held:
             self.held_by.inventory.remove(self)
-            self.map.add_entity(self)
+            self.held_by.map.add_entity(self)
             self.move(self.held_by.y, self.held_by.x)
             self.held = False
             self.held_by = None
@@ -49,7 +49,7 @@ class Item(Entity):
         """
         self.held = True
         self.held_by = player
-        self.map.remove_entity(self)
+        self.held_by.map.remove_entity(self)
         player.inventory.append(self)
 
     def save_state(self) -> dict:
diff --git a/squirrelbattle/interfaces.py b/squirrelbattle/interfaces.py
index 3567ea0..92b4498 100644
--- a/squirrelbattle/interfaces.py
+++ b/squirrelbattle/interfaces.py
@@ -68,7 +68,8 @@ class Map:
         """
         Unregister an entity from the map.
         """
-        self.entities.remove(entity)
+        if entity in self.entities:
+            self.entities.remove(entity)
 
     def find_entities(self, entity_class: type) -> list:
         return [entity for entity in self.entities
diff --git a/squirrelbattle/tests/game_test.py b/squirrelbattle/tests/game_test.py
index 3a32c95..2bb84c5 100644
--- a/squirrelbattle/tests/game_test.py
+++ b/squirrelbattle/tests/game_test.py
@@ -216,6 +216,33 @@ class TestGame(unittest.TestCase):
         self.game.handle_key_pressed(KeyValues.SPACE)
         self.assertEqual(self.game.state, GameMode.MAINMENU)
 
+    def test_mouse_click(self) -> None:
+        """
+        Simulate mouse clicks.
+        """
+        self.game.state = GameMode.MAINMENU
+
+        # Settings menu
+        self.game.display_actions(DisplayActions.MOUSE, 25, 21)
+        self.assertEqual(self.game.main_menu.position, 4)
+        self.assertEqual(self.game.state, GameMode.SETTINGS)
+
+        bomb = Bomb()
+        bomb.hold(self.game.player)
+        bomb2 = Bomb()
+        bomb2.hold(self.game.player)
+
+        self.game.state = GameMode.INVENTORY
+
+        # Click nowhere
+        self.game.display_actions(DisplayActions.MOUSE, 0, 0)
+        self.assertEqual(self.game.state, GameMode.INVENTORY)
+
+        # Click on the second item
+        self.game.display_actions(DisplayActions.MOUSE, 8, 25)
+        self.assertEqual(self.game.state, GameMode.INVENTORY)
+        self.assertEqual(self.game.inventory_menu.position, 1)
+
     def test_new_game(self) -> None:
         """
         Ensure that the start button starts a new game.