Fix BranchDataItemCache (#678)

Author: alexweininger

src/tree/BranchDataItemCache.ts

@@ -3,7 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { isAzExtTreeItem } from '@microsoft/vscode-azext-utils';
 import { ResourceModelBase } from 'api/src/resources/base';
+import { BranchDataItemWrapper } from './BranchDataItemWrapper';
 import { ResourceGroupsItem } from './ResourceGroupsItem';
 
 export class BranchDataItemCache {
@@ -12,8 +14,9 @@ export class BranchDataItemCache {
 
     addBranchItem(branchItem: ResourceModelBase, item: ResourceGroupsItem): void {
         this.branchItemToResourceGroupsItemCache.set(branchItem, item);
-        if (branchItem.id) {
-            this.idToBranchItemCache.set(branchItem.id, branchItem);
+        const id = this.getIdForBranchItem(branchItem);
+        if (id) {
+            this.idToBranchItemCache.set(id, branchItem);
         }
     }
 
@@ -26,11 +29,30 @@ export class BranchDataItemCache {
         return this.branchItemToResourceGroupsItemCache.get(branchItem);
     }
 
-    getItemForId(id?: string): ResourceGroupsItem | undefined {
+    getItemForBranchItemById(branchItem: ResourceModelBase): ResourceGroupsItem | undefined {
+        const id = this.getIdForBranchItem(branchItem);
         if (!id) {
             return undefined;
         }
-        const branchItem = this.idToBranchItemCache.get(id);
-        return branchItem ? this.branchItemToResourceGroupsItemCache.get(branchItem) : undefined;
+        const cachedBranchItem = this.idToBranchItemCache.get(id);
+        return cachedBranchItem ? this.branchItemToResourceGroupsItemCache.get(cachedBranchItem) : undefined;
+    }
+
+    createOrGetItem<T extends BranchDataItemWrapper>(branchItem: ResourceModelBase, createItem: () => T): T {
+        const cachedItem = this.getItemForBranchItemById(branchItem) as T | undefined;
+        if (cachedItem) {
+            cachedItem.branchItem = branchItem;
+            this.addBranchItem(branchItem, cachedItem);
+            return cachedItem;
+        }
+        return createItem();
+    }
+
+    private getIdForBranchItem(branchItem: ResourceModelBase): string | undefined {
+        if (isAzExtTreeItem(branchItem)) {
+            return branchItem.fullId;
+        }
+
+        return branchItem.id;
     }
 }

src/tree/BranchDataItemWrapper.ts

@@ -115,13 +115,9 @@ export class BranchDataItemWrapper implements ResourceGroupsItem, Wrapper {
 export type BranchDataItemFactory = (branchItem: ResourceModelBase, branchDataProvider: BranchDataProvider<ResourceBase, ResourceModelBase>, options?: BranchDataItemOptions) => BranchDataItemWrapper;
 
 export function createBranchDataItemFactory(itemCache: BranchDataItemCache): BranchDataItemFactory {
-    return (branchItem, branchDataProvider, options) => {
-        const cachedItem = itemCache.getItemForId(branchItem.id) as BranchDataItemWrapper | undefined;
-        if (cachedItem) {
-            cachedItem.branchItem = branchItem;
-            itemCache.addBranchItem(branchItem, cachedItem);
-            return cachedItem;
-        }
-        return new BranchDataItemWrapper(branchItem, branchDataProvider, itemCache, options);
-    }
+    return (branchItem, branchDataProvider, options) =>
+        itemCache.createOrGetItem(
+            branchItem,
+            () => new BranchDataItemWrapper(branchItem, branchDataProvider, itemCache, options),
+        )
 }

src/tree/azure/AzureResourceItem.ts

@@ -54,13 +54,9 @@ export class AzureResourceItem<T extends AzureResource> extends BranchDataItemWr
 export type ResourceItemFactory<T extends AzureResource> = (resource: T, branchItem: ResourceModelBase, branchDataProvider: BranchDataProvider<ResourceBase, ResourceModelBase>, parent?: ResourceGroupsItem, options?: BranchDataItemOptions) => AzureResourceItem<T>;
 
 export function createResourceItemFactory<T extends AzureResource>(itemCache: BranchDataItemCache): ResourceItemFactory<T> {
-    return (resource, branchItem, branchDataProvider, parent, options) => {
-        const cachedItem = itemCache.getItemForId(branchItem.id) as AzureResourceItem<T> | undefined;
-        if (cachedItem) {
-            cachedItem.branchItem = branchItem;
-            itemCache.addBranchItem(branchItem, cachedItem);
-            return cachedItem;
-        }
-        return new AzureResourceItem(resource, branchItem, branchDataProvider, itemCache, parent, options);
-    }
+    return (resource, branchItem, branchDataProvider, parent, options) =>
+        itemCache.createOrGetItem(
+            branchItem,
+            () => new AzureResourceItem(resource, branchItem, branchDataProvider, itemCache, parent, options)
+        );
 }