Fix keymapper reverse mapping (#9535)

Fixes #9470

10 files changed, 313 insertions, 74 deletions

diff --git a/all/src/main/templates/release-notes.html b/all/src/main/templates/release-notes.html
index a4fe330c8b..0652dbd58a 100644
--- a/all/src/main/templates/release-notes.html
+++ b/all/src/main/templates/release-notes.html
@@ -107,7 +107,7 @@
         <li>OSGi bundle manifests of Vaadin Framework JARs no longer export <tt>/VAADIN</tt>, and there are new mechanisms for publishing static resources for OSGi</li>
         <li>Tooltip styles for <tt>ContentMode.PREFORMATTED</tt> have been changed in all built-in themes to use the application font and allow long lines to wrap to multiple lines.</li>
         <li><tt>Grid.Column</tt> now extends <tt>AbstractExtension</tt> instead of <tt>AbstractGridExtension</tt> to hide data generator specific API.</li>
-
+        <li><tt>DataCommunicator</tt>, <tt>DataKeyMapper</tt> and <tt>KeyMapper</tt> public APIs have some minor changes for better bean identification.</li>
         <h2>For incompatible or behaviour-altering changes in 8.0, please see <a href="https://vaadin.com/download/release/8.0/8.0.0/release-notes.html#incompatible">8.0 release notes</a></h2>
         
                 
diff --git a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
index 2f46b6526a..0971c36755 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataCommunicator.java
@@ -30,6 +30,7 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import com.vaadin.data.ValueProvider;
 import com.vaadin.data.provider.DataChangeEvent.DataRefreshEvent;
 import com.vaadin.server.AbstractExtension;
 import com.vaadin.server.KeyMapper;
@@ -167,7 +168,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         @Override
         public void generateData(T data, JsonObject jsonObject) {
             // Make sure KeyMapper is up to date
-            getKeyMapper().refresh(data, dataProvider::getId);
+            getKeyMapper().refresh(data);
 
             // Write the key string for given data object
             jsonObject.put(DataCommunicatorConstants.KEY,
@@ -193,7 +194,7 @@ public class DataCommunicator<T> extends AbstractExtension {
     private final ActiveDataHandler handler = new ActiveDataHandler();
 
     /** Empty default data provider. */
-    protected DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
+    private DataProvider<T, ?> dataProvider = new CallbackDataProvider<>(
             q -> Stream.empty(), q -> 0);
     private final DataKeyMapper<T> keyMapper;
 
@@ -212,7 +213,7 @@ public class DataCommunicator<T> extends AbstractExtension {
         addDataGenerator(handler);
         rpc = getRpcProxy(DataCommunicatorClientRpc.class);
         registerRpc(createRpc());
-        keyMapper = createKeyMapper();
+        keyMapper = createKeyMapper(dataProvider::getId);
     }
 
     @Override
@@ -357,13 +358,13 @@ public class DataCommunicator<T> extends AbstractExtension {
 
     /**
      * Fetches a list of items from the DataProvider.
-     * 
+     *
      * @param offset
      *            the starting index of the range
      * @param limit
      *            the max number of results
      * @return the list of items in given range
-     * 
+     *
      * @since 8.1
      */
     @SuppressWarnings({ "rawtypes", "unchecked" })
@@ -585,10 +586,16 @@ public class DataCommunicator<T> extends AbstractExtension {
      * <p>
      * This method is called from the constructor.
      *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
      * @return key mapper
+     *
+     * @since 8.1
+     *
      */
-    protected DataKeyMapper<T> createKeyMapper() {
-        return new KeyMapper<>();
+    protected DataKeyMapper<T> createKeyMapper(ValueProvider<T,Object> identifierGetter) {
+        return new KeyMapper<T>(identifierGetter);
     }
 
     /**
@@ -633,9 +640,7 @@ public class DataCommunicator<T> extends AbstractExtension {
             DataProvider<T, F> dataProvider, F initialFilter) {
         Objects.requireNonNull(dataProvider, "data provider cannot be null");
         filter = initialFilter;
-        detachDataProviderListener();
-        dropAllData();
-        this.dataProvider = dataProvider;
+        setDataProvider(dataProvider);
 
         /*
          * This introduces behavior which influence on the client-server
@@ -672,13 +677,13 @@ public class DataCommunicator<T> extends AbstractExtension {
      * Sets the filter for this DataCommunicator. This method is used by user
      * through the consumer method from {@link #setDataProvider} and should not
      * be called elsewhere.
-     * 
+     *
      * @param filter
      *            the filter
-     * 
+     *
      * @param <F>
      *            the filter type
-     * 
+     *
      * @since 8.1
      */
     protected <F> void setFilter(F filter) {
@@ -725,7 +730,7 @@ public class DataCommunicator<T> extends AbstractExtension {
     /**
      * Getter method for finding the size of DataProvider. Can be overridden by
      * a subclass that uses a specific type of DataProvider and/or query.
-     * 
+     *
      * @return the size of data provider with current filter
      */
     @SuppressWarnings({ "unchecked", "rawtypes" })
@@ -749,7 +754,7 @@ public class DataCommunicator<T> extends AbstractExtension {
                     getUI().access(() -> {
                         if (event instanceof DataRefreshEvent) {
                             T item = ((DataRefreshEvent<T>) event).getItem();
-                            keyMapper.refresh(item, dataProvider::getId);
+                            getKeyMapper().refresh(item);
                             generators.forEach(g -> g.refreshData(item));
                             refresh(item);
                         } else {
@@ -773,4 +778,17 @@ public class DataCommunicator<T> extends AbstractExtension {
             dataProviderUpdateRegistration = null;
         }
     }
+
+    /**
+     * Sets a new {@code DataProvider} and refreshes all the internal structures
+     *
+     * @param dataProvider
+     * @since 8.1
+     */
+    protected void setDataProvider(DataProvider<T, ?> dataProvider) {
+        detachDataProviderListener();
+        dropAllData();
+        this.dataProvider = dataProvider;
+        keyMapper.setIdentifierGetter(dataProvider::getId);
+    }
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
index f5be454f86..6db4a491fe 100644
--- a/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
+++ b/server/src/main/java/com/vaadin/data/provider/DataKeyMapper.java
@@ -79,8 +79,19 @@ public interface DataKeyMapper<T> extends Serializable {
      *
      * @param dataObject
      *            the data object to update
-     * @param identifierGetter
-     *            the function to get an identifier from a data object
+     *
+     * @since 8.1
+     *
+     */
+    void refresh(T dataObject);
+
+    /**
+     * Takes identifier getter into use and updates existing mappings
+     *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
+     * @since 8.1
      */
-    void refresh(T dataObject, ValueProvider<T, Object> identifierGetter);
+    void setIdentifierGetter(ValueProvider<T, Object> identifierGetter);
 }
diff --git a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
index 3e9eb01921..f53287b1df 100644
--- a/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
+++ b/server/src/main/java/com/vaadin/data/provider/HierarchicalDataCommunicator.java
@@ -52,7 +52,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
      */
     public HierarchicalDataCommunicator() {
         super();
-        setDataProvider(new TreeDataProvider<T>(new TreeData<>()), null);
+        setDataProvider(new TreeDataProvider<>(new TreeData<>()), null);
     }
 
     @Override
@@ -220,7 +220,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Returns whether given item has children.
-     * 
+     *
      * @param item
      *            the item to test
      * @return {@code true} if item has children; {@code false} if not
@@ -231,7 +231,7 @@ public class HierarchicalDataCommunicator<T> extends DataCommunicator<T> {
 
     /**
      * Returns whether given item is expanded.
-     * 
+     *
      * @param item
      *            the item to test
      * @return {@code true} if item is expanded; {@code false} if not
diff --git a/server/src/main/java/com/vaadin/server/KeyMapper.java b/server/src/main/java/com/vaadin/server/KeyMapper.java
index 58058c6f69..f7b878dd27 100644
--- a/server/src/main/java/com/vaadin/server/KeyMapper.java
+++ b/server/src/main/java/com/vaadin/server/KeyMapper.java
@@ -18,6 +18,7 @@ package com.vaadin.server;
 
 import java.io.Serializable;
 import java.util.HashMap;
+import java.util.Map;
 
 import com.vaadin.data.ValueProvider;
 import com.vaadin.data.provider.DataKeyMapper;
@@ -33,15 +34,35 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
 
     private int lastKey = 0;
 
-    private final HashMap<V, String> objectKeyMap = new HashMap<>();
+    private final HashMap<Object, String> objectIdKeyMap = new HashMap<>();
 
     private final HashMap<String, V> keyObjectMap = new HashMap<>();
 
+    private ValueProvider<V, Object> identifierGetter;
+
+    /**
+     * Constructs a new mapper
+     *
+     * @param identifierGetter has to return a unique key for every bean, and the returned key has to
+     *                         follow general {@code hashCode()} and {@code equals()} contract,
+     *                         see {@link Object#hashCode()} for details.
+     * @since 8.1
+     */
+    public KeyMapper(ValueProvider<V, Object> identifierGetter) {
+        this.identifierGetter = identifierGetter;
+    }
+
+    /**
+     * Constructs a new mapper with trivial {@code identifierGetter}
+     */
+    public KeyMapper() {
+        this(v -> v);
+    }
+
     /**
      * Gets key for an object.
      *
-     * @param o
-     *            the object.
+     * @param o the object.
      */
     @Override
     public String key(V o) {
@@ -51,14 +72,15 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
         }
 
         // If the object is already mapped, use existing key
-        String key = objectKeyMap.get(o);
+        Object id = identifierGetter.apply(o);
+        String key = objectIdKeyMap.get(id);
         if (key != null) {
             return key;
         }
 
         // If the object is not yet mapped, map it
         key = String.valueOf(++lastKey);
-        objectKeyMap.put(o, key);
+        objectIdKeyMap.put(id, key);
         keyObjectMap.put(key, o);
 
         return key;
@@ -66,14 +88,13 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
 
     @Override
     public boolean has(V o) {
-        return objectKeyMap.containsKey(o);
+        return objectIdKeyMap.containsKey(identifierGetter.apply(o));
     }
 
     /**
      * Retrieves object with the key.
      *
-     * @param key
-     *            the name with the desired value.
+     * @param key the name with the desired value.
      * @return the object with the key.
      */
     @Override
@@ -84,15 +105,12 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
     /**
      * Removes object from the mapper.
      *
-     * @param removeobj
-     *            the object to be removed.
+     * @param removeobj the object to be removed.
      */
     @Override
     public void remove(V removeobj) {
-        final String key = objectKeyMap.get(removeobj);
-
+        final String key = objectIdKeyMap.remove(identifierGetter.apply(removeobj));
         if (key != null) {
-            objectKeyMap.remove(removeobj);
             keyObjectMap.remove(key);
         }
     }
@@ -102,34 +120,39 @@ public class KeyMapper<V> implements DataKeyMapper<V>, Serializable {
      */
     @Override
     public void removeAll() {
-        objectKeyMap.clear();
+        objectIdKeyMap.clear();
         keyObjectMap.clear();
     }
 
     /**
      * Checks if the given key is mapped to an object.
      *
-     * @since 7.7
-     *
-     * @param key
-     *            the key to check
+     * @param key the key to check
      * @return <code>true</code> if the key is currently mapped,
-     *         <code>false</code> otherwise
+     * <code>false</code> otherwise
+     * @since 7.7
      */
     public boolean containsKey(String key) {
         return keyObjectMap.containsKey(key);
     }
 
     @Override
-    public void refresh(V dataObject,
-            ValueProvider<V, Object> identifierGetter) {
+    public void refresh(V dataObject) {
         Object id = identifierGetter.apply(dataObject);
-        objectKeyMap.entrySet().stream()
-                .filter(e -> identifierGetter.apply(e.getKey()).equals(id))
-                .findAny().ifPresent(e -> {
-                    String key = objectKeyMap.remove(e.getKey());
-                    objectKeyMap.put(dataObject, key);
-                    keyObjectMap.put(key, dataObject);
-                });
+        String key = objectIdKeyMap.get(id);
+        if (key != null) {
+            keyObjectMap.put(key, dataObject);
+        }
+    }
+
+    @Override
+    public void setIdentifierGetter(ValueProvider<V, Object> identifierGetter) {
+        if (this.identifierGetter != identifierGetter) {
+            this.identifierGetter = identifierGetter;
+            objectIdKeyMap.clear();
+            for (Map.Entry<String, V> entry : keyObjectMap.entrySet()) {
+                objectIdKeyMap.put(identifierGetter.apply(entry.getValue()),entry.getKey());
+            }
+        }
     }
 }
diff --git a/server/src/main/java/com/vaadin/ui/ComboBox.java b/server/src/main/java/com/vaadin/ui/ComboBox.java
index 62e7ef7751..73d9ad84cf 100644
--- a/server/src/main/java/com/vaadin/ui/ComboBox.java
+++ b/server/src/main/java/com/vaadin/ui/ComboBox.java
@@ -25,6 +25,7 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Stream;
 
+import com.vaadin.data.ValueProvider;
 import org.jsoup.nodes.Element;
 
 import com.vaadin.data.HasFilterableDataProvider;
@@ -188,8 +189,8 @@ public class ComboBox<T> extends AbstractSingleSelect<T>
     public ComboBox() {
         super(new DataCommunicator<T>() {
             @Override
-            protected DataKeyMapper<T> createKeyMapper() {
-                return new KeyMapper<T>() {
+            protected DataKeyMapper<T> createKeyMapper(ValueProvider<T,Object> identifierGetter) {
+                return new KeyMapper<T>(identifierGetter) {
                     @Override
                     public void remove(T removeobj) {
                         // never remove keys from ComboBox to support selection
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);
diff --git a/uitest/src/main/java/com/vaadin/tests/data/GridRefreshWithGetId.java b/uitest/src/main/java/com/vaadin/tests/data/GridRefreshWithGetId.java
new file mode 100644
index 0000000000..b37956d126
--- /dev/null
+++ b/uitest/src/main/java/com/vaadin/tests/data/GridRefreshWithGetId.java
@@ -0,0 +1,109 @@
+package com.vaadin.tests.data;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.vaadin.annotations.Widgetset;
+import com.vaadin.data.provider.ListDataProvider;
+import com.vaadin.server.VaadinRequest;
+import com.vaadin.tests.components.AbstractTestUI;
+import com.vaadin.ui.Button;
+import com.vaadin.ui.Grid;
+
+@Widgetset("com.vaadin.DefaultWidgetSet")
+public class GridRefreshWithGetId extends AbstractTestUI {
+
+    private static class TestObject {
+
+        private final int id;
+        private String name;
+
+        public TestObject(int id, String name) {
+            this.id = id;
+            this.name = name;
+        }
+
+        public int getId() {
+            return id;
+        }
+
+        public String getName() {
+            return name;
+        }
+
+        public void setName(String name) {
+            this.name = name;
+        }
+
+        /**
+         * The class intentionally has strange {@code hashCode()} and {@code equals()}
+         * implementation to ensure if {@code Grid} relies on bean id rather than on
+         * bean hashcode/equals identification.
+         *
+         * {@see Object.equals}
+         */
+        @Override
+        public boolean equals(Object o) {
+            if (this == o)
+                return true;
+            if (o == null || getClass() != o.getClass())
+                return false;
+
+            TestObject myObject = (TestObject) o;
+
+            if (id != myObject.id)
+                return false;
+            return name != null ? name.equals(myObject.name)
+                    : myObject.name == null;
+        }
+
+        /**
+         * The class intentionally has strange {@code hashCode()} and {@code equals()}
+         * implementation to ensure if {@code Grid} relies on bean id rather than on
+         * bean hashcode/equals identification.
+         *
+         * {@see Object.hashCode}
+         */
+        @Override
+        public int hashCode() {
+            int result = id;
+            result = 31 * result + (name != null ? name.hashCode() : 0);
+            return result;
+        }
+    }
+
+    @Override
+    protected void setup(VaadinRequest request) {
+        List<TestObject> data = new ArrayList<>();
+        data.add(new TestObject(0, "blue"));
+        data.add(new TestObject(1, "red"));
+        data.add(new TestObject(2, "green"));
+        data.add(new TestObject(3, "yellow"));
+        data.add(new TestObject(4, "purple"));
+
+        ListDataProvider<TestObject> dataProvider = new ListDataProvider<TestObject>(
+                data) {
+
+            @Override
+            public Object getId(TestObject item) {
+                return item.getId();
+            }
+        };
+
+        Grid<TestObject> grid = new Grid<>();
+        grid.setDataProvider(dataProvider);
+        addComponent(grid);
+
+        grid.addColumn(myObject -> {
+            return myObject.getName();
+        });
+
+        Button button = new Button("Change green to black");
+        button.addClickListener(event1 -> {
+            TestObject myObject = data.get(2);
+            myObject.setName("black");
+            dataProvider.refreshItem(myObject);
+        });
+        addComponent(button);
+    }
+}
diff --git a/uitest/src/test/java/com/vaadin/tests/data/GridRefreshWithGetIdTest.java b/uitest/src/test/java/com/vaadin/tests/data/GridRefreshWithGetIdTest.java
new file mode 100644
index 0000000000..c3bda4f3ca
--- /dev/null
+++ b/uitest/src/test/java/com/vaadin/tests/data/GridRefreshWithGetIdTest.java
@@ -0,0 +1,27 @@
+package com.vaadin.tests.data;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import com.vaadin.testbench.elements.ButtonElement;
+import com.vaadin.testbench.elements.GridElement;
+import com.vaadin.tests.tb3.SingleBrowserTest;
+
+public class GridRefreshWithGetIdTest extends SingleBrowserTest {
+
+    @Test
+    public void testDataIdentifiedAndUpdated() {
+        openTestURL();
+        // Select item
+        GridElement grid = $(GridElement.class).first();
+        grid.getCell(2, 0).click();
+        Assert.assertTrue("Row should be selected",
+                grid.getRow(2).isSelected());
+        Assert.assertEquals("green", grid.getCell(2, 0).getText());
+        $(ButtonElement.class).first().click();
+        Assert.assertTrue("Row was no longer selected",
+                grid.getRow(2).isSelected());
+        Assert.assertEquals("black", grid.getCell(2, 0).getText());
+    }
+
+}