Add AsyncServerTree (#11838)

* wip * Fixes * More fixes * more fixes * Disable when preview features disabled * remove unused imports * Handle promises * PR feedback * Single default ServerGroup color value
2026-01-21 01:25:37 -05:00 · 2020-08-19 14:01:10 -07:00
parent d2e4eeac88
commit 3c538d1c2d
37 changed files with 1654 additions and 506 deletions
--- a/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
+++ b/src/sql/workbench/contrib/objectExplorer/browser/serverTreeView.ts
@@ -38,6 +38,11 @@ import { CommandsRegistry } from 'vs/platform/commands/common/commands';
 import { startsWith } from 'vs/base/common/strings';
 import { SERVER_GROUP_CONFIG } from 'sql/workbench/services/serverGroup/common/interfaces';
 import { horizontalScrollingKey } from 'vs/platform/list/browser/listService';
+import { ITreeContextMenuEvent } from 'vs/base/browser/ui/tree/tree';
+import { ObjectExplorerActionsContext } from 'sql/workbench/services/objectExplorer/browser/objectExplorerActions';
+import { IContextMenuService } from 'vs/platform/contextview/browser/contextView';
+import { IKeybindingService } from 'vs/platform/keybinding/common/keybinding';
+import { AsyncServerTree, ServerTreeElement } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';

 /**
 * ServerTreeview implements the dynamic tree view.
@@ -48,7 +53,7 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 	private _buttonSection: HTMLElement;
 	private _treeSelectionHandler: TreeSelectionHandler;
 	private _activeConnectionsFilterAction: ActiveConnectionsFilterAction;
-	private _tree: ITree;
+	private _tree: ITree | AsyncServerTree;
 	private _onSelectionOrFocusChange: Emitter<void>;
 	private _actionProvider: ServerTreeActionProvider;

@@ -59,7 +64,9 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		@IThemeService private _themeService: IThemeService,
 		@IErrorMessageService private _errorMessageService: IErrorMessageService,
 		@IConfigurationService private _configurationService: IConfigurationService,
-		@ICapabilitiesService capabilitiesService: ICapabilitiesService
+		@ICapabilitiesService capabilitiesService: ICapabilitiesService,
+		@IContextMenuService private _contextMenuService: IContextMenuService,
+		@IKeybindingService private _keybindingService: IKeybindingService
 	) {
 		super();
 		this._activeConnectionsFilterAction = this._instantiationService.createInstance(
@@ -71,9 +78,17 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		this._onSelectionOrFocusChange = new Emitter();
 		this._actionProvider = this._instantiationService.createInstance(ServerTreeActionProvider);
 		capabilitiesService.onCapabilitiesRegistered(async () => {
-			if (this._connectionManagementService.hasRegisteredServers()) {
-				await this.refreshTree();
+			if (this._tree instanceof AsyncServerTree) {
+				// Refresh the tree input now that the capabilities are registered so that we can
+				// get the full ConnectionProfiles with the server info updated properly
+				const treeInput = TreeUpdateUtils.getTreeInput(this._connectionManagementService);
+				await this._tree.setInput(treeInput);
 				this._treeSelectionHandler.onTreeActionStateChange(false);
+			} else {
+				if (this._connectionManagementService.hasRegisteredServers()) {
+					await this.refreshTree();
+					this._treeSelectionHandler.onTreeActionStateChange(false);
+				}
 			}
 		});
 		this.registerCommands();
@@ -97,7 +112,7 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		return this._actionProvider;
 	}

-	public get tree(): ITree {
+	public get tree(): ITree | AsyncServerTree {
 		return this._tree;
 	}

@@ -143,11 +158,30 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		}

 		const horizontalScrollEnabled: boolean = this._configurationService.getValue(horizontalScrollingKey) || false;
-		this._tree = this._register(TreeCreationUtils.createRegisteredServersTree(container, this._instantiationService, horizontalScrollEnabled));
-		//this._tree.setInput(undefined);
+		this._tree = this._register(TreeCreationUtils.createServersTree(container, this._instantiationService, this._configurationService, horizontalScrollEnabled));
 		this._register(this._tree.onDidChangeSelection((event) => this.onSelected(event)));
 		this._register(this._tree.onDidBlur(() => this._onSelectionOrFocusChange.fire()));
 		this._register(this._tree.onDidChangeFocus(() => this._onSelectionOrFocusChange.fire()));
+		if (this._tree instanceof AsyncServerTree) {
+			this._register(this._tree.onContextMenu(e => this.onContextMenu(e)));
+			this._register(this._tree.onMouseDblClick(e => {
+				// Open dashboard on double click for server and database nodes
+				let connectionProfile: ConnectionProfile;
+				if (e.element instanceof ConnectionProfile) {
+					connectionProfile = e.element;
+				} else if (e.element instanceof TreeNode) {
+					if (TreeUpdateUtils.isAvailableDatabaseNode(e.element)) {
+						connectionProfile = TreeUpdateUtils.getConnectionProfile(e.element);
+					}
+				}
+				if (connectionProfile) {
+					this._connectionManagementService.showDashboard(connectionProfile);
+				}
+			}));
+			this._register(this._connectionManagementService.onConnectionChanged(() => {
+				this.refreshTree().catch(err => errors.onUnexpectedError);
+			}));
+		}

 		// Theme styler
 		this._register(attachListStyler(this._tree, this._themeService));
@@ -171,7 +205,7 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 					this.showError(args.errorMessage);
 				}
 				if (args.connection) {
-					this.onObjectExplorerSessionCreated(args.connection);
+					this.onObjectExplorerSessionCreated(args.connection).catch(err => errors.onUnexpectedError);
 				}
 			}));
 		}
@@ -182,7 +216,14 @@ export class ServerTreeView extends Disposable implements IServerTreeView {

 			const expandGroups: boolean = this._configurationService.getValue(SERVER_GROUP_CONFIG)[SERVER_GROUP_AUTOEXPAND_CONFIG];
 			if (expandGroups) {
-				await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(root));
+				if (this._tree instanceof AsyncServerTree) {
+					await Promise.all(ConnectionProfileGroup.getSubgroups(root).map(subgroup => {
+						return this._tree.expand(subgroup);
+					}));
+				} else {
+					await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(root));
+				}
+
 			}

 			if (root && !root.hasValidConnections) {
@@ -200,30 +241,55 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		return uri && startsWith(uri, ConnectionUtils.uriPrefixes.default) && !isBackupRestoreUri;
 	}

-	private async handleAddConnectionProfile(newProfile: IConnectionProfile): Promise<void> {
-		if (newProfile) {
-			const groups = this._connectionManagementService.getConnectionGroups();
-			const profile = ConnectionUtils.findProfileInGroup(newProfile, groups);
-			if (profile) {
-				newProfile = profile;
-			}
-		}
-
+	private async handleAddConnectionProfile(newProfile?: IConnectionProfile): Promise<void> {
 		if (this._buttonSection) {
 			hide(this._buttonSection);
 			this._activeConnectionsFilterAction.enabled = true;
 		}
-		const currentSelections = this._tree.getSelection();
-		const currentSelectedElement = currentSelections && currentSelections.length >= 1 ? currentSelections[0] : undefined;
-		const newProfileIsSelected = currentSelectedElement && newProfile ? currentSelectedElement.id === newProfile.id : false;
-		if (newProfile && currentSelectedElement && !newProfileIsSelected) {
-			this._tree.clearSelection();
-		}
-		await this.refreshTree();
-		if (newProfile && !newProfileIsSelected) {
-			await this._tree.reveal(newProfile);
-			this._tree.select(newProfile);
+
+		if (this._tree instanceof AsyncServerTree) {
+			// When new connection groups are added the event is fired with undefined so
+			// we still want to refresh the tree in that case to pick up the changes
+			await this.refreshTree();
+			if (newProfile) {
+				const currentSelections = this._tree.getSelection();
+				const currentSelectedElement = currentSelections && currentSelections.length >= 1 ? currentSelections[0] : undefined;
+				const newProfileIsSelected = currentSelectedElement && currentSelectedElement.id === newProfile.id;
+				// Clear any other selected elements first
+				if (currentSelectedElement && !newProfileIsSelected) {
+					this._tree.setSelection([]);
+				}
+				const newConnectionProfile = this.getConnectionInTreeInput(newProfile.id);
+				if (newConnectionProfile) {
+					// Re-render to update the connection status badge
+					this._tree.rerender(newConnectionProfile);
+					this._tree.setSelection([newConnectionProfile]);
+					this._tree.expand(newConnectionProfile);
+				}
+			}
+
+		} else {
+			if (newProfile) {
+				const groups = this._connectionManagementService.getConnectionGroups();
+				const profile = ConnectionUtils.findProfileInGroup(newProfile, groups);
+				if (profile) {
+					newProfile = profile;
+				}
+			}
+
+			const currentSelections = this._tree.getSelection();
+			const currentSelectedElement = currentSelections && currentSelections.length >= 1 ? currentSelections[0] : undefined;
+			const newProfileIsSelected = currentSelectedElement && newProfile ? currentSelectedElement.id === newProfile.id : false;
+			if (newProfile && currentSelectedElement && !newProfileIsSelected) {
+				this._tree.clearSelection();
+			}
+			await this.refreshTree();
+			if (newProfile && !newProfileIsSelected) {
+				await this._tree.reveal(newProfile);
+				this._tree.select(newProfile);
+			}
 		}
+
 	}

 	private showError(errorMessage: string) {
@@ -232,66 +298,83 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		}
 	}

-	private getConnectionInTreeInput(connectionId: string): ConnectionProfile {
-		const root = TreeUpdateUtils.getTreeInput(this._connectionManagementService);
-		const connections = ConnectionProfileGroup.getConnectionsInGroup(root);
-		const results = connections.filter(con => {
-			if (connectionId === con.id) {
-				return true;
-			} else {
-				return false;
+	/**
+	 * Gets the ConnectionProfile object in the tree for the specified ID, or undefined if it doesn't exist.
+	 * @param connectionId The connection ID to search for
+	 */
+	private getConnectionInTreeInput(connectionId: string): ConnectionProfile | undefined {
+		if (this._tree instanceof AsyncServerTree) {
+			const root = this._tree.getInput();
+			const connections = ConnectionProfileGroup.getConnectionsInGroup(root);
+			return connections.find(conn => conn.id === connectionId);
+		} else {
+			const root = TreeUpdateUtils.getTreeInput(this._connectionManagementService);
+			const connections = ConnectionProfileGroup.getConnectionsInGroup(root);
+			const results = connections.filter(con => {
+				if (connectionId === con.id) {
+					return true;
+				} else {
+					return false;
+				}
+			});
+			if (results && results.length > 0) {
+				return results[0];
 			}
-		});
-		if (results && results.length > 0) {
-			return results[0];
+			return undefined;
 		}
-		return null;
 	}

-	private onObjectExplorerSessionCreated(connection: IConnectionProfile) {
-		const conn = this.getConnectionInTreeInput(connection.id);
-		if (conn) {
-			this._tree.refresh(conn).then(() => {
-				return this._tree.expand(conn).then(() => {
-					return this._tree.reveal(conn, 0.5).then(() => {
-						this._treeSelectionHandler.onTreeActionStateChange(false);
-					});
-				});
-			}).then(null, errors.onUnexpectedError);
+	private async onObjectExplorerSessionCreated(connection: IConnectionProfile): Promise<void> {
+		const element = this.getConnectionInTreeInput(connection.id);
+		if (element) {
+			if (this._tree instanceof AsyncServerTree) {
+				this._tree.rerender(element);
+			} else {
+				await this._tree.refresh(element);
+			}
+			await this._tree.expand(element);
+			await this._tree.reveal(element, 0.5);
+			this._treeSelectionHandler.onTreeActionStateChange(false);
 		}
 	}

 	public addObjectExplorerNodeAndRefreshTree(connection: IConnectionProfile): void {
 		hide(this.messages);
 		if (!this._objectExplorerService.getObjectExplorerNode(connection)) {
-			this._objectExplorerService.updateObjectExplorerNodes(connection).then(() => {
-				// The oe request is sent. an event will be raised when the session is created
-			}, error => {
-			});
+			this._objectExplorerService.updateObjectExplorerNodes(connection).catch(e => errors.onUnexpectedError(e));
 		}
 	}

-	public deleteObjectExplorerNodeAndRefreshTree(connection: IConnectionProfile): Promise<void> {
+	public async deleteObjectExplorerNodeAndRefreshTree(connection: IConnectionProfile): Promise<void> {
 		if (connection) {
 			const conn = this.getConnectionInTreeInput(connection.id);
 			if (conn) {
-				return this._objectExplorerService.deleteObjectExplorerNode(conn).then(async () => {
+				await this._objectExplorerService.deleteObjectExplorerNode(conn);
+				if (this._tree instanceof AsyncServerTree) {
+					// Collapse the node before refreshing so the refresh doesn't try to fetch
+					// the children again (which causes it to try and connect)
+					this._tree.collapse(conn);
+					await this.refreshTree();
+				} else {
 					await this._tree.collapse(conn);
 					return this._tree.refresh(conn);
-				});
+				}
 			}
 		}
-		return Promise.resolve();
 	}

-	public refreshTree(): Promise<void> {
+	public async refreshTree(): Promise<void> {
 		hide(this.messages);
 		this.clearOtherActions();
 		return TreeUpdateUtils.registeredServerUpdate(this._tree, this._connectionManagementService);
 	}

-	public refreshElement(element: any): Promise<void> {
-		return this._tree.refresh(element);
+	public async refreshElement(element: ServerTreeElement): Promise<void> {
+		if (this._tree instanceof AsyncServerTree) {
+			return this._tree.updateChildren(element);
+		} else {
+			return this._tree.refresh(element);
+		}
 	}

 	/**
@@ -350,9 +433,19 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 			this._tree.setInput(treeInput).then(async () => {
 				if (isHidden(this.messages)) {
 					this._tree.getFocus();
-					await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+					if (this._tree instanceof AsyncServerTree) {
+						await Promise.all(ConnectionProfileGroup.getSubgroups(treeInput).map(subgroup => {
+							this._tree.expand(subgroup);
+						}));
+					} else {
+						await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+					}
 				} else {
-					this._tree.clearFocus();
+					if (this._tree instanceof AsyncServerTree) {
+						this._tree.setFocus([]);
+					} else {
+						this._tree.clearFocus();
+					}
 				}
 			}, errors.onUnexpectedError);
 		} else {
@@ -382,9 +475,19 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		this._tree.setInput(treeInput).then(async () => {
 			if (isHidden(this.messages)) {
 				this._tree.getFocus();
-				await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+				if (this._tree instanceof AsyncServerTree) {
+					await Promise.all(ConnectionProfileGroup.getSubgroups(treeInput).map(subgroup => {
+						this._tree.expand(subgroup);
+					}));
+				} else {
+					await this._tree.expandAll(ConnectionProfileGroup.getSubgroups(treeInput));
+				}
 			} else {
-				this._tree.clearFocus();
+				if (this._tree instanceof AsyncServerTree) {
+					this._tree.setFocus([]);
+				} else {
+					this._tree.clearFocus();
+				}
 			}
 		}, errors.onUnexpectedError);
 	}
@@ -450,17 +553,6 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 		this._tree.layout(height);
 	}

-	/**
-	 * set the visibility of the view
-	 */
-	public setVisible(visible: boolean): void {
-		if (visible) {
-			this._tree.onVisible();
-		} else {
-			this._tree.onHidden();
-		}
-	}
-
 	/**
 	 * Get the list of selected nodes in the tree
 	*/
@@ -472,48 +564,108 @@ export class ServerTreeView extends Disposable implements IServerTreeView {
 	 * Get whether the tree view currently has focus
 	*/
 	public isFocused(): boolean {
-		return this._tree.isDOMFocused();
+		return this._tree.getHTMLElement() === document.activeElement;
 	}

 	/**
 	 * Set whether the given element is expanded or collapsed
 	 */
-	public setExpandedState(element: TreeNode | ConnectionProfile, expandedState: TreeItemCollapsibleState): Promise<void> {
+	public async setExpandedState(element: ServerTreeElement, expandedState: TreeItemCollapsibleState): Promise<void> {
 		if (expandedState === TreeItemCollapsibleState.Collapsed) {
 			return this._tree.collapse(element);
 		} else if (expandedState === TreeItemCollapsibleState.Expanded) {
 			return this._tree.expand(element);
 		}
-		return Promise.resolve();
 	}

 	/**
 	 * Reveal the given element in the tree
 	 */
-	public reveal(element: TreeNode | ConnectionProfile): Promise<void> {
+	public async reveal(element: ServerTreeElement): Promise<void> {
 		return this._tree.reveal(element);
 	}

 	/**
 	 * Select the given element in the tree and clear any other selections
 	 */
-	public setSelected(element: TreeNode | ConnectionProfile, selected: boolean, clearOtherSelections: boolean): Promise<void> {
+	public async setSelected(element: ServerTreeElement, selected: boolean, clearOtherSelections: boolean): Promise<void> {
 		if (clearOtherSelections || (selected && clearOtherSelections !== false)) {
-			this._tree.clearSelection();
+			if (this._tree instanceof AsyncServerTree) {
+				this._tree.setSelection([]);
+			} else {
+				this._tree.clearSelection();
+			}
+
 		}
 		if (selected) {
-			this._tree.select(element);
-			return this._tree.reveal(element);
+			if (this._tree instanceof AsyncServerTree) {
+				this._tree.setSelection(this._tree.getSelection().concat(element));
+				this._tree.reveal(element);
+			} else {
+				this._tree.select(element);
+				return this._tree.reveal(element);
+			}
 		} else {
-			this._tree.deselect(element);
-			return Promise.resolve();
+			if (this._tree instanceof AsyncServerTree) {
+				this._tree.setSelection(this._tree.getSelection().filter(item => item !== element));
+			} else {
+				this._tree.deselect(element);
+			}
 		}
 	}

 	/**
 	 * Check if the given element in the tree is expanded
 	 */
-	public isExpanded(element: TreeNode | ConnectionProfile): boolean {
-		return this._tree.isExpanded(element);
+	public isExpanded(element: ServerTreeElement): boolean {
+		if (this._tree instanceof AsyncServerTree) {
+			return !this._tree.getNode(element).collapsed;
+		} else {
+			return this._tree.isExpanded(element);
+		}
+
+	}
+
+	/**
+	 * Return actions in the context menu
+	 */
+	private onContextMenu(e: ITreeContextMenuEvent<ServerTreeElement>): boolean {
+		e.browserEvent.preventDefault();
+		e.browserEvent.stopPropagation();
+		this._tree.setSelection([e.element]);
+
+		let actionContext: any;
+		if (e.element instanceof TreeNode) {
+			let context = new ObjectExplorerActionsContext();
+			context.nodeInfo = e.element.toNodeInfo();
+			// Note: getting DB name before, but intentionally not using treeUpdateUtils.getConnectionProfile as it replaces
+			// the connection ID with a new one. This breaks a number of internal tasks
+			context.connectionProfile = e.element.getConnectionProfile().toIConnectionProfile();
+			context.connectionProfile.databaseName = e.element.getDatabaseName();
+			actionContext = context;
+		} else if (e.element instanceof ConnectionProfile) {
+			let context = new ObjectExplorerActionsContext();
+			context.connectionProfile = e.element.toIConnectionProfile();
+			context.isConnectionNode = true;
+			actionContext = context;
+		} else {
+			// TODO: because the connection group is used as a context object and isn't serializable,
+			// the Group-level context menu is not currently extensible
+			actionContext = e.element;
+		}
+
+		this._contextMenuService.showContextMenu({
+			getAnchor: () => e.anchor,
+			getActions: () => this._actionProvider.getActions(this._tree, e.element),
+			getKeyBinding: (action) => this._keybindingService.lookupKeybinding(action.id),
+			onHide: (wasCancelled?: boolean) => {
+				if (wasCancelled) {
+					this._tree.domFocus();
+				}
+			},
+			getActionsContext: () => (actionContext)
+		});
+
+		return true;
 	}
 }
--- a/src/sql/workbench/contrib/objectExplorer/common/serverGroup.contribution.ts
+++ b/src/sql/workbench/contrib/objectExplorer/common/serverGroup.contribution.ts
@@ -8,6 +8,7 @@ import { Registry } from 'vs/platform/registry/common/platform';
 import { IJSONSchema } from 'vs/base/common/jsonSchema';
 import { localize } from 'vs/nls';
 import { SERVER_GROUP_CONFIG, SERVER_GROUP_COLORS_CONFIG } from 'sql/workbench/services/serverGroup/common/interfaces';
+import { DefaultServerGroupColor } from 'sql/workbench/services/serverGroup/common/serverGroupViewModel';

 const configurationRegistry = Registry.as<IConfigurationRegistry>(Extensions.Configuration);

@@ -29,7 +30,7 @@ const serverGroupConfig: IConfigurationNode = {
 				'#98AFC7',
 				'#4452A6',
 				'#6A6599',
-				'#515151'
+				DefaultServerGroupColor
 			]
 		},
 		[SERVER_GROUP_CONFIG + '.' + SERVER_GROUP_AUTOEXPAND_CONFIG]: {
@@ -40,4 +41,18 @@ const serverGroupConfig: IConfigurationNode = {
 	}
 };

+const serverTreeConfig: IConfigurationNode = {
+	'id': 'serverTree',
+	'title': 'Server Tree',
+	'type': 'object',
+	'properties': {
+		'serverTree.useAsyncServerTree': {
+			'type': 'boolean',
+			'default': true,
+			'description': localize('serverTree.useAsyncServerTree', "(Preview) Use the new async server tree for the Servers view and Connection Dialog with support for new features such as dynamic node filtering.")
+		}
+	}
+};
+
 configurationRegistry.registerConfiguration(serverGroupConfig);
+configurationRegistry.registerConfiguration(serverTreeConfig);
--- a/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts
+++ b/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts
@@ -20,8 +20,6 @@ import * as  LocalizedConstants from 'sql/workbench/services/connection/browser/
 import { ObjectExplorerService, ObjectExplorerNodeEventArgs } from 'sql/workbench/services/objectExplorer/browser/objectExplorerService';
 import { TreeNode } from 'sql/workbench/services/objectExplorer/common/treeNode';
 import { NodeType } from 'sql/workbench/services/objectExplorer/common/nodeType';
-import { Tree } from 'vs/base/parts/tree/browser/treeImpl';
-import { ServerTreeDataSource } from 'sql/workbench/services/objectExplorer/browser/serverTreeDataSource';
 import { Emitter, Event } from 'vs/base/common/event';
 import Severity from 'vs/base/common/severity';
 import { ObjectExplorerActionsContext } from 'sql/workbench/services/objectExplorer/browser/objectExplorerActions';
@@ -35,6 +33,12 @@ import { IViewsService, IView, ViewContainerLocation, ViewContainer, IViewPaneCo
 import { ConsoleLogService } from 'vs/platform/log/common/log';
 import { IProgressIndicator } from 'vs/platform/progress/common/progress';
 import { IPaneComposite } from 'vs/workbench/common/panecomposite';
+import { MockContextKeyService } from 'vs/platform/keybinding/test/common/mockKeybindingService';
+import { TestAccessibilityService, TestListService } from 'vs/workbench/test/browser/workbenchTestServices';
+import { TestConfigurationService } from 'sql/platform/connection/test/common/testConfigurationService';
+import { ServerTreeDataSource } from 'sql/workbench/services/objectExplorer/browser/serverTreeDataSource';
+import { Tree } from 'vs/base/parts/tree/browser/treeImpl';
+import { AsyncServerTree } from 'sql/workbench/services/objectExplorer/browser/asyncServerTree';

 suite('SQL Connection Tree Action tests', () => {
 	let errorMessageService: TypeMoq.Mock<TestErrorMessageService>;
@@ -152,8 +156,14 @@ suite('SQL Connection Tree Action tests', () => {
 			}
 		};

-		let manageConnectionAction: OEManageConnectionAction = new OEManageConnectionAction(OEManageConnectionAction.ID,
-			OEManageConnectionAction.LABEL, connectionManagementService.object, capabilitiesService, instantiationService.object, objectExplorerService.object, viewsService);
+		let manageConnectionAction: OEManageConnectionAction = new OEManageConnectionAction(
+			OEManageConnectionAction.ID,
+			OEManageConnectionAction.LABEL,
+			connectionManagementService.object,
+			capabilitiesService,
+			instantiationService.object,
+			objectExplorerService.object,
+			viewsService);

 		let actionContext = new ObjectExplorerActionsContext();
 		actionContext.connectionProfile = connection.toIConnectionProfile();
@@ -190,8 +200,14 @@ suite('SQL Connection Tree Action tests', () => {
 			return treeSelectionMock.object;
 		});

-		let manageConnectionAction: OEManageConnectionAction = new OEManageConnectionAction(OEManageConnectionAction.ID,
-			OEManageConnectionAction.LABEL, connectionManagementService.object, capabilitiesService, instantiationService.object, objectExplorerService.object, undefined);
+		let manageConnectionAction: OEManageConnectionAction = new OEManageConnectionAction(
+			OEManageConnectionAction.ID,
+			OEManageConnectionAction.LABEL,
+			connectionManagementService.object,
+			capabilitiesService,
+			instantiationService.object,
+			objectExplorerService.object,
+			undefined);

 		let actionContext = new ObjectExplorerActionsContext();
 		actionContext.connectionProfile = connection.toIConnectionProfile();
@@ -444,7 +460,7 @@ suite('SQL Connection Tree Action tests', () => {
 		tree.setup(x => x.refresh(TypeMoq.It.isAny())).returns(() => Promise.resolve(null));
 		tree.setup(x => x.expand(TypeMoq.It.isAny())).returns(() => Promise.resolve(null));
 		tree.setup(x => x.collapse(TypeMoq.It.isAny())).returns(() => Promise.resolve(null));
-		let connectionAction: RefreshAction = new RefreshAction(RefreshAction.ID,
+		let refreshAction: RefreshAction = new RefreshAction(RefreshAction.ID,
 			RefreshAction.LABEL,
 			tree.object,
 			connection,
@@ -453,7 +469,7 @@ suite('SQL Connection Tree Action tests', () => {
 			undefined,
 			logService);

-		return connectionAction.run().then((value) => {
+		return refreshAction.run().then((value) => {
 			connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
 			objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
 			objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
@@ -547,6 +563,200 @@ suite('SQL Connection Tree Action tests', () => {
 		});
 	});

+	// chgagnon TODO - skipping for now since mocking and instanceof don't work well together. Will re-enable once old tree is removed
+	test.skip('RefreshConnectionAction - AsyncServerTree - refresh should be called if connection status is connect', () => {
+		let isConnectedReturnValue: boolean = true;
+		let sqlProvider = {
+			providerId: mssqlProviderName,
+			displayName: 'MSSQL',
+			connectionOptions: [],
+		};
+
+		capabilitiesService.capabilities[mssqlProviderName] = { connection: sqlProvider };
+
+		let connection = new ConnectionProfile(capabilitiesService, {
+			connectionName: 'Test',
+			savePassword: false,
+			groupFullName: 'testGroup',
+			serverName: 'testServerName',
+			databaseName: 'testDatabaseName',
+			authenticationType: 'inetgrated',
+			password: 'test',
+			userName: 'testUsername',
+			groupId: undefined,
+			providerName: mssqlProviderName,
+			options: {},
+			saveProfile: true,
+			id: 'testID'
+		});
+		let conProfGroup = new ConnectionProfileGroup('testGroup', undefined, 'testGroup', undefined, undefined);
+		conProfGroup.connections = [connection];
+		let connectionManagementService = TypeMoq.Mock.ofType(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+		connectionManagementService.callBase = true;
+		connectionManagementService.setup(x => x.getConnectionGroups()).returns(() => [conProfGroup]);
+		connectionManagementService.setup(x => x.getActiveConnections()).returns(() => [connection]);
+		connectionManagementService.setup(x => x.addSavedPassword(TypeMoq.It.isAny())).returns(() => new Promise<ConnectionProfile>((resolve) => {
+			resolve(connection);
+		}));
+		connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue);
+
+		let objectExplorerSession = {
+			success: true,
+			sessionId: '1234',
+			rootNode: {
+				nodePath: 'testServerName\tables',
+				nodeType: NodeType.Folder,
+				label: 'Tables',
+				isLeaf: false,
+				metadata: null,
+				nodeSubType: '',
+				nodeStatus: '',
+				errorMessage: ''
+			},
+			errorMessage: ''
+		};
+
+		let tablesNode = new TreeNode(NodeType.Folder, 'Tables', false, 'testServerName\Db1\tables', '', '', null, null, undefined, undefined);
+		tablesNode.connection = connection;
+		tablesNode.session = objectExplorerSession;
+		let table1Node = new TreeNode(NodeType.Table, 'dbo.Table1', false, 'testServerName\tables\dbo.Table1', '', '', tablesNode, null, undefined, undefined);
+		let table2Node = new TreeNode(NodeType.Table, 'dbo.Table1', false, 'testServerName\tables\dbo.Table1', '', '', tablesNode, null, undefined, undefined);
+		tablesNode.children = [table1Node, table2Node];
+		let objectExplorerService = TypeMoq.Mock.ofType(ObjectExplorerService, TypeMoq.MockBehavior.Loose, connectionManagementService.object);
+		objectExplorerService.callBase = true;
+		objectExplorerService.setup(x => x.getObjectExplorerNode(TypeMoq.It.isAny())).returns(() => tablesNode);
+		objectExplorerService.setup(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve([table1Node, table2Node]));
+		const treeMock = TypeMoq.Mock.ofType(AsyncServerTree, TypeMoq.MockBehavior.Strict,
+			'ConnectionTreeActionsTest', // user
+			$('div'), // container
+			{}, // delegate
+			[], // renderers
+			{}, // data source
+			{}, // options
+			new MockContextKeyService(), // IContextKeyService
+			new TestListService(), // IListService,
+			undefined, // IThemeService,
+			new TestConfigurationService(), // IConfigurationService,
+			undefined, // IKeybindingService,
+			new TestAccessibilityService()); // IAccessibilityService
+		treeMock.callBase = true;
+		treeMock.setup(x => x.expand(TypeMoq.It.isAny())).returns(() => Promise.resolve(null));
+		treeMock.setup(x => x.collapse(TypeMoq.It.isAny())).returns(() => true);
+		treeMock.setup(x => x.updateChildren(TypeMoq.It.isAny())).returns(() => Promise.resolve());
+		let refreshAction: RefreshAction = new RefreshAction(RefreshAction.ID,
+			RefreshAction.LABEL,
+			treeMock.instance,
+			connection,
+			connectionManagementService.object,
+			objectExplorerService.object,
+			undefined,
+			logService);
+
+		return refreshAction.run().then((value) => {
+			connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
+			objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
+			objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
+			treeMock.verify(x => x.updateChildren(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
+		});
+	});
+
+	test('RefreshConnectionAction - AsyncServerTree - refresh should not be called if connection status is not connect', () => {
+		let isConnectedReturnValue: boolean = false;
+		let sqlProvider = {
+			providerId: mssqlProviderName,
+			displayName: 'MSSQL',
+			connectionOptions: []
+		};
+
+		capabilitiesService.capabilities[mssqlProviderName] = { connection: sqlProvider };
+
+		let connection = new ConnectionProfile(capabilitiesService, {
+			connectionName: 'Test',
+			savePassword: false,
+			groupFullName: 'testGroup',
+			serverName: 'testServerName',
+			databaseName: 'testDatabaseName',
+			authenticationType: 'inetgrated',
+			password: 'test',
+			userName: 'testUsername',
+			groupId: undefined,
+			providerName: mssqlProviderName,
+			options: {},
+			saveProfile: true,
+			id: 'testID'
+		});
+		let conProfGroup = new ConnectionProfileGroup('testGroup', undefined, 'testGroup', undefined, undefined);
+		conProfGroup.connections = [connection];
+		let connectionManagementService = TypeMoq.Mock.ofType(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
+		connectionManagementService.callBase = true;
+		connectionManagementService.setup(x => x.getConnectionGroups()).returns(() => [conProfGroup]);
+		connectionManagementService.setup(x => x.getActiveConnections()).returns(() => [connection]);
+		connectionManagementService.setup(x => x.addSavedPassword(TypeMoq.It.isAny())).returns(() => new Promise<ConnectionProfile>((resolve) => {
+			resolve(connection);
+		}));
+		connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue);
+
+		let objectExplorerSession = {
+			success: true,
+			sessionId: '1234',
+			rootNode: {
+				nodePath: 'testServerName\tables',
+				nodeType: NodeType.Folder,
+				label: 'Tables',
+				isLeaf: false,
+				metadata: null,
+				nodeSubType: '',
+				nodeStatus: '',
+				errorMessage: ''
+			},
+			errorMessage: ''
+		};
+
+		let tablesNode = new TreeNode(NodeType.Folder, 'Tables', false, 'testServerName\Db1\tables', '', '', null, null, undefined, undefined);
+		tablesNode.connection = connection;
+		tablesNode.session = objectExplorerSession;
+		let table1Node = new TreeNode(NodeType.Table, 'dbo.Table1', false, 'testServerName\tables\dbo.Table1', '', '', tablesNode, null, undefined, undefined);
+		let table2Node = new TreeNode(NodeType.Table, 'dbo.Table1', false, 'testServerName\tables\dbo.Table1', '', '', tablesNode, null, undefined, undefined);
+		tablesNode.children = [table1Node, table2Node];
+		let objectExplorerService = TypeMoq.Mock.ofType(ObjectExplorerService, TypeMoq.MockBehavior.Loose, connectionManagementService.object);
+		objectExplorerService.callBase = true;
+		objectExplorerService.setup(x => x.getObjectExplorerNode(TypeMoq.It.isAny())).returns(() => tablesNode);
+		objectExplorerService.setup(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve([table1Node, table2Node]));
+		let tree = TypeMoq.Mock.ofType<AsyncServerTree>(AsyncServerTree, TypeMoq.MockBehavior.Loose,
+			'ConnectionTreeActionsTest', // user
+			$('div'), // container
+			{}, // delegate
+			[], // renderers
+			{}, // data source
+			{}, // options
+			new MockContextKeyService(), // IContextKeyService
+			new TestListService(), // IListService,
+			undefined, // IThemeService,
+			new TestConfigurationService(), // IConfigurationService,
+			undefined, // IKeybindingService,
+			new TestAccessibilityService()); // IAccessibilityService
+		tree.callBase = true;
+
+		tree.setup(x => x.updateChildren(TypeMoq.It.isAny())).returns(() => Promise.resolve());
+		tree.setup(x => x.expand(TypeMoq.It.isAny())).returns(() => Promise.resolve(null));
+		let connectionAction: RefreshAction = new RefreshAction(RefreshAction.ID,
+			RefreshAction.LABEL,
+			tree.object,
+			connection,
+			connectionManagementService.object,
+			objectExplorerService.object,
+			undefined,
+			logService);
+
+		return connectionAction.run().then((value) => {
+			connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
+			objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
+			objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
+			tree.verify(x => x.updateChildren(TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
+			tree.verify(x => x.expand(TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
+		});
+	});
+
 	test('EditConnectionAction - test if show connection dialog is called', () => {
 		let connectionManagementService = createConnectionManagementService(true, undefined);

--- a/src/sql/workbench/contrib/objectExplorer/test/browser/serverTreeView.test.ts
+++ b/src/sql/workbench/contrib/objectExplorer/test/browser/serverTreeView.test.ts
@@ -39,7 +39,7 @@ suite('ServerTreeView onAddConnectionProfile handler tests', () => {
 		);
 		mockConnectionManagementService.setup(x => x.getConnectionGroups()).returns(x => []);
 		mockConnectionManagementService.setup(x => x.hasRegisteredServers()).returns(() => true);
-		serverTreeView = new ServerTreeView(mockConnectionManagementService.object, instantiationService, undefined, new TestThemeService(), undefined, undefined, capabilitiesService);
+		serverTreeView = new ServerTreeView(mockConnectionManagementService.object, instantiationService, undefined, new TestThemeService(), undefined, undefined, capabilitiesService, undefined, undefined);
 		mockTree = TypeMoq.Mock.ofType(TestTree);
 		(serverTreeView as any)._tree = mockTree.object;
 		mockRefreshTreeMethod = TypeMoq.Mock.ofType(Function);
@@ -64,15 +64,6 @@ suite('ServerTreeView onAddConnectionProfile handler tests', () => {
 		mockTree.verify(x => x.layout(TypeMoq.It.isAnyNumber()), TypeMoq.Times.once());
 	});

-	test('setVisibility', async () => {
-		mockTree.setup(x => x.onVisible());
-		mockTree.setup(x => x.onHidden());
-		serverTreeView.setVisible(true);
-		mockTree.verify(x => x.onVisible(), TypeMoq.Times.once());
-		serverTreeView.setVisible(false);
-		mockTree.verify(x => x.onHidden(), TypeMoq.Times.once());
-	});
-
 	test('getSelection', async () => {
 		mockTree.setup(x => x.getSelection());

@@ -81,9 +72,9 @@ suite('ServerTreeView onAddConnectionProfile handler tests', () => {
 	});

 	test('isFocused', async () => {
-		mockTree.setup(x => x.isDOMFocused());
+		mockTree.setup(x => x.getHTMLElement());
 		serverTreeView.isFocused();
-		mockTree.verify(x => x.isDOMFocused(), TypeMoq.Times.once());
+		mockTree.verify(x => x.getHTMLElement(), TypeMoq.Times.once());
 	});

 	test('reveal', async () => {