From d5ad8cd4efb971779122c15f5a24e375dcf7717e Mon Sep 17 00:00:00 2001
From: Matt Seddon <37993418+mattseddon@users.noreply.github.com>
Date: Tue, 2 May 2023 07:08:52 +1000
Subject: [PATCH] Display progress errors to users (#3791)

---
 extension/src/setup/autoInstall.ts            | 42 +++++++------
 .../src/test/suite/setup/autoInstall.test.ts  | 60 +++++++++++++++++++
 extension/src/vscode/toast.ts                 | 16 +++++
 3 files changed, 101 insertions(+), 17 deletions(-)

diff --git a/extension/src/setup/autoInstall.ts b/extension/src/setup/autoInstall.ts
index 49b4b43421..b38480e43a 100644
--- a/extension/src/setup/autoInstall.ts
+++ b/extension/src/setup/autoInstall.ts
@@ -23,25 +23,33 @@ const showInstallProgress = (
   Toast.showProgress('Installing packages', async progress => {
     progress.report({ increment: 0 })
 
-    await Toast.runCommandAndIncrementProgress(
-      async () => {
-        await installPackages(root, pythonBinPath, 'dvclive')
-        return 'DVCLive Installed'
-      },
-      progress,
-      25
-    )
+    try {
+      await Toast.runCommandAndIncrementProgress(
+        async () => {
+          await installPackages(root, pythonBinPath, 'dvclive')
+          return 'DVCLive Installed'
+        },
+        progress,
+        25
+      )
+    } catch (error: unknown) {
+      return Toast.reportProgressError(error, progress)
+    }
 
-    await Toast.runCommandAndIncrementProgress(
-      async () => {
-        await installPackages(root, pythonBinPath, 'dvc')
-        return 'DVC Installed'
-      },
-      progress,
-      75
-    )
+    try {
+      await Toast.runCommandAndIncrementProgress(
+        async () => {
+          await installPackages(root, pythonBinPath, 'dvc')
+          return 'DVC Installed'
+        },
+        progress,
+        75
+      )
 
-    return Toast.delayProgressClosing()
+      return Toast.delayProgressClosing()
+    } catch (error: unknown) {
+      return Toast.reportProgressError(error, progress)
+    }
   })
 
 export const autoInstallDvc = async (): Promise<unknown> => {
diff --git a/extension/src/test/suite/setup/autoInstall.test.ts b/extension/src/test/suite/setup/autoInstall.test.ts
index 44276f5dfd..775a75fb32 100644
--- a/extension/src/test/suite/setup/autoInstall.test.ts
+++ b/extension/src/test/suite/setup/autoInstall.test.ts
@@ -8,6 +8,7 @@ import * as Python from '../../../python'
 import { autoInstallDvc } from '../../../setup/autoInstall'
 import * as WorkspaceFolders from '../../../vscode/workspaceFolders'
 import { bypassProgressCloseDelay } from '../util'
+import { Toast } from '../../../vscode/toast'
 
 const { getDefaultPython } = Python
 
@@ -89,5 +90,64 @@ suite('Auto Install Test Suite', () => {
         'dvclive'
       )
     })
+
+    it('should show an error message and exit early if DVCLive fails to install', async () => {
+      bypassProgressCloseDelay()
+      const cwd = __dirname
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(defaultPython)
+      const mockInstallPackages = stub(Python, 'installPackages').rejects(
+        new Error('Failed to install DVCLive')
+      )
+      stub(WorkspaceFolders, 'getFirstWorkspaceFolder').returns(cwd)
+
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+
+      await autoInstallDvc()
+
+      expect(showProgressSpy).to.be.called
+      expect(showErrorSpy).not.to.be.called
+      expect(mockInstallPackages).to.be.calledOnce
+      expect(mockInstallPackages).to.be.calledWithExactly(
+        cwd,
+        defaultPython,
+        'dvclive'
+      )
+    })
+
+    it('should show an error message if DVC fails to install', async () => {
+      bypassProgressCloseDelay()
+      const cwd = __dirname
+      stub(PythonExtension, 'getPythonExecutionDetails').resolves(undefined)
+      stub(Python, 'findPythonBin').resolves(defaultPython)
+      const mockInstallPackages = stub(Python, 'installPackages')
+        .onFirstCall()
+        .resolves(undefined)
+        .onSecondCall()
+        .rejects(new Error('Network error'))
+      stub(WorkspaceFolders, 'getFirstWorkspaceFolder').returns(cwd)
+
+      const showProgressSpy = spy(window, 'withProgress')
+      const showErrorSpy = spy(window, 'showErrorMessage')
+      const reportProgressErrorSpy = spy(Toast, 'reportProgressError')
+
+      await autoInstallDvc()
+
+      expect(showProgressSpy).to.be.called
+      expect(showErrorSpy).not.to.be.called
+      expect(reportProgressErrorSpy).to.be.calledOnce
+      expect(mockInstallPackages).to.be.calledTwice
+      expect(mockInstallPackages).to.be.calledWithExactly(
+        cwd,
+        defaultPython,
+        'dvclive'
+      )
+      expect(mockInstallPackages).to.be.calledWithExactly(
+        cwd,
+        defaultPython,
+        'dvc'
+      )
+    })
   })
 })
diff --git a/extension/src/vscode/toast.ts b/extension/src/vscode/toast.ts
index d8006950a2..cd4c81579e 100644
--- a/extension/src/vscode/toast.ts
+++ b/extension/src/vscode/toast.ts
@@ -75,6 +75,22 @@ export class Toast {
     })
   }
 
+  static reportProgressError(
+    error: unknown,
+    progress: Progress<{
+      message?: string | undefined
+      increment?: number | undefined
+    }>
+  ) {
+    const message = (error as Error)?.message || 'an unexpected error occurred'
+
+    progress.report({
+      increment: 0,
+      message
+    })
+    return Toast.delayProgressClosing(60000)
+  }
+
   static delayProgressClosing(ms = 5000) {
     return delay(ms)
   }