fix(menu): recompute index before marking selection (#5047)

packages/mdc-menu/foundation.ts

@@ -97,8 +97,10 @@ export class MDCMenuFoundation extends MDCFoundation<MDCMenuAdapter> {
 
     // Wait for the menu to close before adding/removing classes that affect styles.
     this.closeAnimationEndTimerId_ = setTimeout(() => {
-      if (this.adapter_.isSelectableItemAtIndex(index)) {
-        this.setSelectedIndex(index);
+      // Recompute the index in case the menu contents have changed.
+      const recomputedIndex = this.adapter_.getElementIndex(listItem);
+      if (this.adapter_.isSelectableItemAtIndex(recomputedIndex)) {
+        this.setSelectedIndex(recomputedIndex);
       }
     }, MDCMenuSurfaceFoundation.numbers.TRANSITION_CLOSE_DURATION);
   }

test/unit/mdc-menu/menu.foundation.test.js

@@ -225,6 +225,30 @@ test('handleItemAction item action event inside of a child element of a selectio
     {times: 0});
 });
 
+test('handleItemAction adds class to the correct child element of a selection group when menu has mutated', () => {
+  const {foundation, mockAdapter, clock} = setupTest();
+  const itemEl = document.createElement('li');
+  td.when(mockAdapter.elementContainsClass(itemEl, listClasses.LIST_ITEM_CLASS)).thenReturn(true);
+  td.when(mockAdapter.getElementIndex(itemEl)).thenReturn(1);
+  td.when(mockAdapter.elementContainsClass(itemEl, cssClasses.MENU_SELECTION_GROUP)).thenReturn(true);
+
+  td.when(mockAdapter.isSelectableItemAtIndex(1)).thenReturn(true);
+  td.when(mockAdapter.getSelectedSiblingOfItemAtIndex(1)).thenReturn(-1);
+  td.when(mockAdapter.getMenuItemCount()).thenReturn(2);
+
+  foundation.handleItemAction(itemEl);
+
+  // Element at index 1 is now at index 0
+  td.when(mockAdapter.getElementIndex(itemEl)).thenReturn(0);
+  td.when(mockAdapter.isSelectableItemAtIndex(0)).thenReturn(true);
+  td.when(mockAdapter.getSelectedSiblingOfItemAtIndex(0)).thenReturn(-1);
+  td.when(mockAdapter.getMenuItemCount()).thenReturn(1);
+
+  clock.tick(numbers.TRANSITION_CLOSE_DURATION);
+
+  td.verify(mockAdapter.addClassToElementAtIndex(0, cssClasses.MENU_SELECTED_LIST_ITEM), {times: 1});
+});
+
 test('handleMenuSurfaceOpened menu focuses the list root element by default on menu surface opened', () => {
   const {foundation, mockAdapter} = setupTest();