await on expand

Fixes microsoft#229878

Author: andreamah

src/vs/workbench/contrib/search/browser/searchActionsRemoveReplace.ts

@@ -126,7 +126,7 @@ registerAction2(class RemoveAction extends Action2 {
 
 		if (focusElement && shouldRefocusMatch) {
 			if (!nextFocusElement) {
-				nextFocusElement = getLastNodeFromSameType(viewer, focusElement);
+				nextFocusElement = await getLastNodeFromSameType(viewer, focusElement);
 			}
 
 			if (nextFocusElement && !arrayContainsElementOrParent(nextFocusElement, elementsToRemove)) {
@@ -294,7 +294,7 @@ async function performReplace(accessor: ServicesAccessor,
 
 	if (focusElement) {
 		if (!nextFocusElement) {
-			nextFocusElement = getLastNodeFromSameType(viewer, focusElement);
+			nextFocusElement = await getLastNodeFromSameType(viewer, focusElement);
 		}
 
 		if (nextFocusElement) {
@@ -389,13 +389,16 @@ export function getElementToFocusAfterRemoved(viewer: WorkbenchCompressibleAsync
 /***
  * Finds the last element in the tree with the same type as `element`
  */
-export function getLastNodeFromSameType(viewer: WorkbenchCompressibleAsyncDataTree<SearchResult, RenderableMatch>, element: RenderableMatch): RenderableMatch | undefined {
+export async function getLastNodeFromSameType(viewer: WorkbenchCompressibleAsyncDataTree<SearchResult, RenderableMatch>, element: RenderableMatch): Promise<RenderableMatch | undefined> {
 	let lastElem: RenderableMatch | null = viewer.lastVisibleElement ?? null;
 
 	while (lastElem) {
 		const compareVal = compareLevels(element, lastElem);
 		if (compareVal === -1) {
-			viewer.expand(lastElem);
+			const expanded = await viewer.expand(lastElem);
+			if (!expanded) {
+				return lastElem;
+			}
 			lastElem = viewer.lastVisibleElement;
 		} else if (compareVal === 1) {
 			const potentialLastElem = viewer.getParentElement(lastElem);

src/vs/workbench/contrib/search/test/browser/searchActions.test.ts

@@ -75,25 +75,25 @@ suite('Search Actions', () => {
 		assert.strictEqual(data[4], actual);
 	});
 
-	test('Find last FileMatch in Tree', function () {
+	test('Find last FileMatch in Tree', async function () {
 		const fileMatch1 = aFileMatch();
 		const fileMatch2 = aFileMatch();
 		const fileMatch3 = aFileMatch();
 		const data = [fileMatch1, aMatch(fileMatch1), fileMatch2, aMatch(fileMatch2), fileMatch3, aMatch(fileMatch3)];
 		const tree = aTree(data);
 
-		const actual = getLastNodeFromSameType(tree, fileMatch1);
+		const actual = await getLastNodeFromSameType(tree, fileMatch1);
 		assert.strictEqual(fileMatch3, actual);
 	});
 
-	test('Find last Match in Tree', function () {
+	test('Find last Match in Tree', async function () {
 		const fileMatch1 = aFileMatch();
 		const fileMatch2 = aFileMatch();
 		const fileMatch3 = aFileMatch();
 		const data = [fileMatch1, aMatch(fileMatch1), fileMatch2, aMatch(fileMatch2), fileMatch3, aMatch(fileMatch3)];
 		const tree = aTree(data);
 
-		const actual = getLastNodeFromSameType(tree, aMatch(fileMatch1));
+		const actual = await getLastNodeFromSameType(tree, aMatch(fileMatch1));
 		assert.strictEqual(data[5], actual);
 	});