Fix keymapper reverse mapping (#9535)

Fixes #9470

2 files changed, 72 insertions, 22 deletions

diff --git a/server/src/test/java/com/vaadin/tests/server/KeyMapperIdGetterTest.java b/server/src/test/java/com/vaadin/tests/server/KeyMapperIdGetterTest.java
new file mode 100644
index 0000000000..de51b5cffc
--- /dev/null
+++ b/server/src/test/java/com/vaadin/tests/server/KeyMapperIdGetterTest.java
@@ -0,0 +1,42 @@
+package com.vaadin.tests.server;
+
+import com.vaadin.server.KeyMapper;
+
+/**
+ * The test checks the same functionality as {@link KeyMapperTest} does, but uses custom {@code identifierGetter}
+ * instead of default trivial one. {@code BrokenBean} intentionally has broken {@code hashCode} and {@code equals}, and the test should pass
+ * despite of that, because {@code BrokenBean.getId()} is used for bean identification.
+ */
+public class KeyMapperIdGetterTest extends KeyMapperTest {
+
+
+    private static class BrokenBean {
+        private final Object id = new Object();
+
+        @Override
+        public int hashCode() {
+            return 0;
+        }
+
+        @Override
+        public boolean equals(Object obj) {
+            return false;
+        }
+
+        public Object getId() {
+            return id;
+        }
+    }
+
+    protected Object createObject() {
+        return new BrokenBean();
+    }
+
+    protected KeyMapper<Object> createKeyMapper() {
+
+        KeyMapper<BrokenBean> keyMapper = new KeyMapper<>();
+        keyMapper.setIdentifierGetter(BrokenBean::getId);
+        return (KeyMapper) keyMapper;
+    }
+
+}
diff --git a/server/src/test/java/com/vaadin/tests/server/KeyMapperTest.java b/server/src/test/java/com/vaadin/tests/server/KeyMapperTest.java
index 5b882ac7b3..47a81aa8e4 100644
--- a/server/src/test/java/com/vaadin/tests/server/KeyMapperTest.java
+++ b/server/src/test/java/com/vaadin/tests/server/KeyMapperTest.java
@@ -1,33 +1,33 @@
 package com.vaadin.tests.server;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotSame;
-import static org.junit.Assert.fail;
+import com.vaadin.server.KeyMapper;
+import org.junit.Test;
 
 import java.lang.reflect.Field;
 import java.util.HashMap;
 
-import org.junit.Test;
-
-import com.vaadin.server.KeyMapper;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.fail;
 
 public class KeyMapperTest {
 
     @Test
     public void testAdd() {
-        KeyMapper<Object> mapper = new KeyMapper<>();
-        Object o1 = new Object();
-        Object o2 = new Object();
-        Object o3 = new Object();
+        KeyMapper<Object> mapper = createKeyMapper();
+        Object o1 = createObject();
+        Object o2 = createObject();
+        Object o3 = createObject();
 
         // Create new ids
         String key1 = mapper.key(o1);
         String key2 = mapper.key(o2);
         String key3 = mapper.key(o3);
 
-        assertEquals(mapper.get(key1), o1);
-        assertEquals(mapper.get(key2), o2);
-        assertEquals(mapper.get(key3), o3);
+        assertSame(mapper.get(key1), o1);
+        assertSame(mapper.get(key2), o2);
+        assertSame(mapper.get(key3), o3);
         assertNotSame(key1, key2);
         assertNotSame(key1, key3);
         assertNotSame(key2, key3);
@@ -45,12 +45,20 @@ public class KeyMapperTest {
 
     }
 
+    protected Object createObject() {
+        return new Object();
+    }
+
+    protected KeyMapper<Object> createKeyMapper() {
+        return new KeyMapper<>();
+    }
+
     @Test
     public void testRemoveAll() {
-        KeyMapper<Object> mapper = new KeyMapper<>();
-        Object o1 = new Object();
-        Object o2 = new Object();
-        Object o3 = new Object();
+        KeyMapper<Object> mapper = createKeyMapper();
+        Object o1 = createObject();
+        Object o2 = createObject();
+        Object o3 = createObject();
 
         // Create new ids
         mapper.key(o1);
@@ -65,10 +73,10 @@ public class KeyMapperTest {
 
     @Test
     public void testRemove() {
-        KeyMapper<Object> mapper = new KeyMapper<>();
-        Object o1 = new Object();
-        Object o2 = new Object();
-        Object o3 = new Object();
+        KeyMapper<Object> mapper = createKeyMapper();
+        Object o1 = createObject();
+        Object o2 = createObject();
+        Object o3 = createObject();
 
         // Create new ids
         mapper.key(o1);
@@ -91,7 +99,7 @@ public class KeyMapperTest {
 
     private void assertSize(KeyMapper<?> mapper, int i) {
         try {
-            Field f1 = KeyMapper.class.getDeclaredField("objectKeyMap");
+            Field f1 = KeyMapper.class.getDeclaredField("objectIdKeyMap");
             Field f2 = KeyMapper.class.getDeclaredField("keyObjectMap");
             f1.setAccessible(true);
             f2.setAccessible(true);