Fix PRs in PRs tree collapses on refresh (microsoft#5634)

* Fix PRs in PRs tree collapses on refresh
Fixes microsoft#5556

* Fix test

src/test/view/prsTree.test.ts

@@ -29,6 +29,7 @@ import { LoggingOctokit, RateLogger } from '../../github/loggingOctokit';
 import { GitHubServerType } from '../../common/authentication';
 import { DataUri } from '../../common/uri';
 import { IAccount, ITeam } from '../../github/interface';
+import { asPromise } from '../../common/utils';
 
 describe('GitHub Pull Requests view', function () {
 	let sinon: SinonSandbox;
@@ -196,6 +197,7 @@ describe('GitHub Pull Requests view', function () {
 
 			// Need to call getChildren twice to get past the quick render with an empty list
 			await localNode!.getChildren();
+			await asPromise(provider.prsTreeModel.onLoaded);
 			const localChildren = await localNode!.getChildren();
 			assert.strictEqual(localChildren.length, 2);
 			const [localItem0, localItem1] = await Promise.all(localChildren.map(node => node.getTreeItem()));

src/view/prsTreeDataProvider.ts

@@ -41,22 +41,22 @@ export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<Tre
 	private _view: vscode.TreeView<TreeNode>;
 	private _initialized: boolean = false;
 	public notificationProvider: NotificationProvider;
-	private _prsTreeModel: PrsTreeModel;
+	public readonly prsTreeModel: PrsTreeModel;
 
 	get view(): vscode.TreeView<TreeNode> {
 		return this._view;
 	}
 
 	constructor(private _telemetry: ITelemetry, private readonly _context: vscode.ExtensionContext, private readonly _reposManager: RepositoriesManager) {
 		this._disposables = [];
-		this._prsTreeModel = new PrsTreeModel(this._telemetry, this._reposManager, _context);
-		this._disposables.push(this._prsTreeModel);
-		this._disposables.push(this._prsTreeModel.onDidChangeData(folderManager => folderManager ? this.refreshRepo(folderManager) : this.refresh()));
-		this._disposables.push(new PRStatusDecorationProvider(this._prsTreeModel));
+		this.prsTreeModel = new PrsTreeModel(this._telemetry, this._reposManager, _context);
+		this._disposables.push(this.prsTreeModel);
+		this._disposables.push(this.prsTreeModel.onDidChangeData(folderManager => folderManager ? this.refreshRepo(folderManager) : this.refresh()));
+		this._disposables.push(new PRStatusDecorationProvider(this.prsTreeModel));
 		this._disposables.push(vscode.window.registerFileDecorationProvider(DecorationProvider));
 		this._disposables.push(
 			vscode.commands.registerCommand('pr.refreshList', _ => {
-				this._prsTreeModel.clearCache();
+				this.prsTreeModel.clearCache();
 				this._onDidChangeTreeData.fire();
 			}),
 		);
@@ -111,10 +111,10 @@ export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<Tre
 		this._disposables.push(this._view.onDidChangeCheckboxState(TreeUtils.processCheckboxUpdates));
 
 		this._disposables.push(this._view.onDidExpandElement(expanded => {
-			this._prsTreeModel.updateExpandedQueries(expanded.element, true);
+			this.prsTreeModel.updateExpandedQueries(expanded.element, true);
 		}));
 		this._disposables.push(this._view.onDidCollapseElement(collapsed => {
-			this._prsTreeModel.updateExpandedQueries(collapsed.element, false);
+			this.prsTreeModel.updateExpandedQueries(collapsed.element, false);
 		}));
 	}
 
@@ -263,7 +263,7 @@ export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<Tre
 					this,
 					this.notificationProvider,
 					this._context,
-					this._prsTreeModel
+					this.prsTreeModel
 				);
 			} else {
 				result = this._reposManager.folderManagers.map(
@@ -275,7 +275,7 @@ export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<Tre
 							this._telemetry,
 							this.notificationProvider,
 							this._context,
-							this._prsTreeModel
+							this.prsTreeModel
 						),
 				);
 			}

src/view/prsTreeModel.ts

@@ -38,6 +38,9 @@ export class PrsTreeModel implements vscode.Disposable {
 	private readonly _onDidChangeData: vscode.EventEmitter<FolderRepositoryManager | void> = new vscode.EventEmitter();
 	public readonly onDidChangeData = this._onDidChangeData.event;
 	private _expandedQueries: Set<string> = new Set();
+	private _hasLoaded: boolean = false;
+	private _onLoaded: vscode.EventEmitter<void> = new vscode.EventEmitter();
+	public readonly onLoaded = this._onLoaded.event;
 
 	// Key is identifier from createPRNodeUri
 	private readonly _queriedPullRequests: Map<string, PRStatusChange> = new Map();
@@ -95,6 +98,15 @@ export class PrsTreeModel implements vscode.Disposable {
 		return this._expandedQueries;
 	}
 
+	get hasLoaded(): boolean {
+		return this._hasLoaded;
+	}
+
+	private set hasLoaded(value: boolean) {
+		this._hasLoaded = value;
+		this._onLoaded.fire();
+	}
+
 	public cachedPRStatus(identifier: string): PRStatusChange | undefined {
 		return this._queriedPullRequests.get(identifier);
 	}
@@ -188,6 +200,7 @@ export class PrsTreeModel implements vscode.Disposable {
 		this._telemetry.sendTelemetryEvent('pr.expand.local');
 		// Don't await this._getChecks. It fires an event that will be listened to.
 		this._getChecks(prs);
+		this.hasLoaded = true;
 		return { hasMorePages: false, hasUnsearchedRepositories: false, items: prs };
 	}
 
@@ -210,6 +223,7 @@ export class PrsTreeModel implements vscode.Disposable {
 		this._telemetry.sendTelemetryEvent('pr.expand.query');
 		// Don't await this._getChecks. It fires an event that will be listened to.
 		this._getChecks(prs.items);
+		this.hasLoaded = true;
 		return prs;
 	}
 
@@ -231,6 +245,7 @@ export class PrsTreeModel implements vscode.Disposable {
 		this._telemetry.sendTelemetryEvent('pr.expand.all');
 		// Don't await this._getChecks. It fires an event that will be listened to.
 		this._getChecks(prs.items);
+		this.hasLoaded = true;
 		return prs;
 	}
 

src/view/treeNodes/categoryNode.ts

@@ -124,7 +124,6 @@ export class CategoryTreeNode extends TreeNode implements vscode.TreeItem {
 	public repositoryPageInformation: Map<string, PageInformation> = new Map<string, PageInformation>();
 	public contextValue: string;
 	public readonly id: string = '';
-	private _firstLoad: boolean = true;
 
 	constructor(
 		public parent: TreeNodeParent,
@@ -306,8 +305,7 @@ export class CategoryTreeNode extends TreeNode implements vscode.TreeItem {
 
 	async getChildren(): Promise<TreeNode[]> {
 		await super.getChildren();
-		if (this._firstLoad) {
-			this._firstLoad = false;
+		if (!this._prsTreeModel.hasLoaded) {
 			this.doGetChildren().then(() => this.refresh(this));
 			return [];
 		}