Add progress indicator when fetching a studio token

extension/src/setup/studio.ts

@@ -8,6 +8,7 @@ import { ContextKey, setContextValue } from '../vscode/context'
 import { Disposable } from '../class/dispose'
 import { getCallBackUrl, openUrl, waitForUriResponse } from '../vscode/external'
 import { Modal } from '../vscode/modal'
+import { Toast } from '../vscode/toast'
 
 export const isStudioAccessToken = (text?: string): boolean => {
   if (!text) {
@@ -86,7 +87,7 @@ export class Studio extends Disposable {
     const storedToken = this.getStudioAccessToken()
     const isConnected = isStudioAccessToken(storedToken)
     this.studioIsConnected = isConnected
-    return setContextValue(ContextKey.STUDIO_CONNECTED, isConnected)
+    await setContextValue(ContextKey.STUDIO_CONNECTED, isConnected)
   }
 
   public async updateStudioOffline(shareLive: boolean) {
@@ -169,15 +170,27 @@ export class Studio extends Disposable {
     return accessToken
   }
 
-  private async requestStudioToken(deviceCode: string, tokenUri: string) {
-    const token = await this.fetchStudioToken(deviceCode, tokenUri)
-    const cwd = this.getCwd()
+  private requestStudioToken(deviceCode: string, tokenUri: string) {
+    return Toast.showProgress('Connecting to Studio', async progress => {
+      progress.report({ increment: 0 })
+      progress.report({ increment: 25, message: 'Fetching token...' })
+      const token = await this.fetchStudioToken(deviceCode, tokenUri)
+      const cwd = this.getCwd()
 
-    if (!token || !cwd) {
-      return
-    }
+      if (!token || !cwd) {
+        const error = new Error('Connection failed')
+        return Toast.reportProgressError(error, progress)
+      }
 
-    return this.saveStudioAccessTokenInConfig(cwd, token)
+      await this.saveStudioAccessTokenInConfig(cwd, token)
+
+      progress.report({
+        increment: 75,
+        message: 'Token saved'
+      })
+
+      return Toast.delayProgressClosing(15000)
+    })
   }
 
   private async setStudioValues() {

extension/src/test/suite/setup/index.test.ts

@@ -17,7 +17,8 @@ import {
   closeAllEditors,
   getMessageReceivedEmitter,
   quickPickInitialized,
-  selectQuickPickItem
+  selectQuickPickItem,
+  waitForSpyCall
 } from '../util'
 import { WEBVIEW_TEST_TIMEOUT } from '../timeouts'
 import { MessageFromWebviewType } from '../../../webview/contract'
@@ -53,6 +54,7 @@ import { DvcConfig } from '../../../cli/dvc/config'
 import * as QuickPickUtil from '../../../setup/quickPick'
 import { EventName } from '../../../telemetry/constants'
 import { Modal } from '../../../vscode/modal'
+import { Toast } from '../../../vscode/toast'
 
 suite('Setup Test Suite', () => {
   const disposable = Disposable.fn()
@@ -923,22 +925,18 @@ suite('Setup Test Suite', () => {
         `${mockStudioRes.verification_uri}?redirect_uri=${mockCallbackUrl}&code=${mockStudioRes.user_code}`
       )
 
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any
-      const mockGetCwd = stub(studio as any, 'getCwd')
       const mockModalShowError = stub(Modal, 'errorWithOptions')
       const mockSaveStudioToken = stub(studio, 'saveStudioAccessTokenInConfig')
+      const showProgressSpy = spy(Toast, 'showProgress')
+      const reportProgressErrorSpy = spy(Toast, 'reportProgressError')
+      const delayProgressClosingSpy = spy(Toast, 'delayProgressClosing')
       mockFetch.onSecondCall().resolves({
         json: () =>
           Promise.resolve({ detail: 'Request failed for some reason.' }),
         status: 500
       } as Fetch.Response)
 
-      const failedTokenEvent = new Promise(resolve =>
-        mockGetCwd.onFirstCall().callsFake(() => {
-          resolve(undefined)
-          return dvcDemoPath
-        })
-      )
+      const failedTokenEvent = waitForSpyCall(delayProgressClosingSpy, 0)
 
       mockOnStudioResponse()
 
@@ -958,19 +956,16 @@ suite('Setup Test Suite', () => {
       expect(mockModalShowError).to.be.calledWithExactly(
         'Unable to get token. Failed with "Request failed for some reason."'
       )
+      expect(showProgressSpy).to.be.calledOnce
+      expect(reportProgressErrorSpy).to.be.calledOnce
       expect(mockSaveStudioToken).not.to.be.called
 
       const mockToken = 'isat_AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'
       mockFetch.onThirdCall().resolves({
         json: () => Promise.resolve({ access_token: mockToken }),
         status: 200
       } as Fetch.Response)
-      const tokenEvent = new Promise(resolve =>
-        mockGetCwd.onSecondCall().callsFake(() => {
-          resolve(undefined)
-          return dvcDemoPath
-        })
-      )
+      const tokenEvent = waitForSpyCall(delayProgressClosingSpy, 1)
 
       mockOnStudioResponse()
 
@@ -982,6 +977,9 @@ suite('Setup Test Suite', () => {
         dvcDemoPath,
         mockToken
       )
+      expect(showProgressSpy).to.be.calledTwice
+      expect(reportProgressErrorSpy).not.to.be.calledTwice
+      expect(delayProgressClosingSpy).to.be.calledTwice
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it("should handle a message from the webview to manually save the user's Studio access token", async () => {

webview/src/setup/components/App.test.tsx

@@ -631,7 +631,7 @@ describe('App', () => {
   })
 
   describe('Studio not connected', () => {
-    it('should show buttons which request a token from Studio or add an already creatd one', () => {
+    it('should show buttons which request a token from Studio or add an already created one', () => {
       renderApp()
       mockPostMessage.mockClear()
       const getTokenButton = screen.getByText('Get Token')

webview/src/setup/components/studio/Connect.tsx

@@ -7,7 +7,7 @@ import { Button } from '../../../shared/components/button/Button'
 export const Connect: React.FC = () => {
   return (
     <EmptyState isFullScreen={false}>
-      <div data-testid="setup-studio-content">
+      <div>
         <h1>
           Connect to <a href={STUDIO_URL}>Studio</a>
         </h1>