Ensure DVC tracked tree is rebuilt when a dvc yaml error has been fixed (#3300)

Author: mattseddon

extension/src/repository/index.ts

@@ -96,8 +96,6 @@ export class Repository extends DeferredDisposable {
     this.treeDataChanged.fire()
     this.sourceControlManagement.setState(sourceControlManagementState)
     this.scmDecorationProvider.setState(scmDecorationState)
-    if (errorDecorationState) {
-      this.errorDecorationProvider.setState(errorDecorationState)
-    }
+    this.errorDecorationProvider.setState(errorDecorationState)
   }
 }

extension/src/repository/model/decorationProvider.ts

@@ -1,6 +1,7 @@
 import { EventEmitter, FileDecoration, Uri } from 'vscode'
 import { DecoratableTreeItemScheme, getDecoratableUri } from '../../tree'
 import { ErrorDecorationProvider } from '../../tree/errorDecorationProvider'
+import { uniqueValues } from '../../util/array'
 
 export class DecorationProvider extends ErrorDecorationProvider {
   private errors = new Set<string>()
@@ -15,10 +16,10 @@ export class DecorationProvider extends ErrorDecorationProvider {
     }
   }
 
-  public setState(errors: Set<string>) {
+  public setState(errors: Set<string> = new Set()) {
     const urisToUpdate: Uri[] = []
 
-    for (const label of errors) {
+    for (const label of uniqueValues([...errors, ...this.errors])) {
       urisToUpdate.push(getDecoratableUri(label, this.scheme))
     }
     this.errors = errors

extension/src/repository/model/index.test.ts

@@ -231,5 +231,34 @@ describe('RepositoryModel', () => {
         untracked: []
       })
     })
+
+    it('should clear an error when there is no data in the tree and the error is fixed', () => {
+      const emptyRepoData = {
+        dataStatus: {},
+        hasGitChanges: true,
+        untracked: new Set<string>()
+      }
+
+      const msg = "'./dvc.yaml' validation failed.\n\nwith some other text"
+
+      const error = {
+        msg,
+        type: 'caught error'
+      }
+
+      const emptyRepoError = { ...emptyRepoData, dataStatus: { error } }
+      const model = new RepositoryModel(dvcDemoPath)
+
+      model.transformAndSet(emptyRepoData)
+      expect(model.getChildren(dvcDemoPath)).toStrictEqual([])
+
+      model.transformAndSet(emptyRepoError)
+      expect(model.getChildren(dvcDemoPath)).toStrictEqual([
+        { error: { label: './dvc.yaml validation failed.', msg } }
+      ])
+
+      model.transformAndSet(emptyRepoData)
+      expect(model.getChildren(dvcDemoPath)).toStrictEqual([])
+    })
   })
 })

extension/src/repository/model/index.ts

@@ -87,7 +87,11 @@ export class RepositoryModel extends Disposable {
   }
 
   private collectTree(tracked: Set<string>): void {
-    if (!sameContents([...tracked], [...this.getTracked()])) {
+    const errorNeedsCleared = this.tree.get(this.dvcRoot)?.[0]?.error
+    if (
+      errorNeedsCleared ||
+      !sameContents([...tracked], [...this.getTracked()])
+    ) {
       this.tracked = tracked
       this.tree = collectTree(this.dvcRoot, this.tracked)
     }