Switch add Studio access token to update when Studio is connected

iterative · Mar 6, 2023 · 855c278 · 855c278
1 parent e3bdb7d
commit 855c278
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 2 deletions.
diff --git a/extension/package.json b/extension/package.json
@@ -117,6 +117,11 @@
         "command": "dvc.addStudioAccessToken",
         "category": "DVC"
       },
+      {
+        "title": "Update Studio Access Token",
+        "command": "dvc.updateStudioAccessToken",
+        "category": "DVC"
+      },
       {
         "title": "Add",
         "command": "dvc.addTarget",
@@ -635,7 +640,11 @@
         },
         {
           "command": "dvc.addStudioAccessToken",
-          "when": "dvc.commands.available && dvc.project.available"
+          "when": "dvc.commands.available && dvc.project.available && !dvc.studio.connected"
+        },
+        {
+          "command": "dvc.updateStudioAccessToken",
+          "when": "dvc.commands.available && dvc.project.available && dvc.studio.connected"
         },
         {
           "command": "dvc.addTarget",

diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts
@@ -99,6 +99,7 @@ export enum RegisteredCommands {
   CONNECT_SHOW = 'dvc.showConnect',
   OPEN_STUDIO_SETTINGS = 'dvc.showStudioSettings',
   ADD_STUDIO_ACCESS_TOKEN = 'dvc.addStudioAccessToken',
+  UPDATE_STUDIO_ACCESS_TOKEN = 'dvc.updateStudioAccessToken',
   REMOVE_STUDIO_ACCESS_TOKEN = 'dvc.removeStudioAccessToken',
   EXPERIMENT_VIEW_SHARE_TO_STUDIO = 'dvc.views.experiments.shareExperimentToStudio'
 }
diff --git a/extension/src/connect/register.ts b/extension/src/connect/register.ts
@@ -21,6 +21,11 @@ export const registerConnectCommands = (
     () => connect.saveStudioAccessToken()
   )
 
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.UPDATE_STUDIO_ACCESS_TOKEN,
+    () => connect.saveStudioAccessToken()
+  )
+
   internalCommands.registerExternalCommand(
     RegisteredCommands.REMOVE_STUDIO_ACCESS_TOKEN,
     () => connect.removeStudioAccessToken()

diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts
@@ -287,5 +287,6 @@ export interface IEventNamePropertyMapping {
   [EventName.CONNECT_SHOW]: undefined
   [EventName.OPEN_STUDIO_SETTINGS]: undefined
   [EventName.ADD_STUDIO_ACCESS_TOKEN]: undefined
+  [EventName.UPDATE_STUDIO_ACCESS_TOKEN]: undefined
   [EventName.REMOVE_STUDIO_ACCESS_TOKEN]: undefined
 }
diff --git a/webview/src/connect/components/App.test.tsx b/webview/src/connect/components/App.test.tsx
@@ -96,5 +96,16 @@ describe('App', () => {
         type: MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE
       })
     })
+
+    it('should enable the user to update their studio token', () => {
+      const shareExperimentsLive = false
+      renderApp(true, shareExperimentsLive)
+      mockPostMessage.mockClear()
+      const button = screen.getByText('Update Token')
+      fireEvent.click(button)
+      expect(mockPostMessage).toHaveBeenCalledWith({
+        type: MessageFromWebviewType.SAVE_STUDIO_TOKEN
+      })
+    })
   })
 })
diff --git a/webview/src/connect/components/Studio.tsx b/webview/src/connect/components/Studio.tsx
@@ -82,7 +82,7 @@ const Settings: React.FC<{
           appearance="primary"
           isNested={false}
           text={'Update Token'}
-          onClick={openStudio}
+          onClick={saveStudioToken}
         />
         <Button
           appearance="secondary"