Transfer selection/starred status to renamed experiment (#4783)

Author: mattseddon

extension/src/experiments/commands/index.ts

@@ -19,8 +19,19 @@ export const getBranchExperimentCommand =
 
 export const getRenameExperimentCommand =
   (experiments: WorkspaceExperiments) =>
-  (cwd: string, oldName: string, newName: string) =>
-    experiments.runCommand(AvailableCommands.EXP_RENAME, cwd, oldName, newName)
+  async (cwd: string, oldName: string, newName: string) => {
+    const output = await experiments.runCommand(
+      AvailableCommands.EXP_RENAME,
+      cwd,
+      oldName,
+      newName
+    )
+    if (!output) {
+      return
+    }
+    const repository = experiments.getRepository(cwd)
+    return repository.transferExperimentDetails(oldName, newName)
+  }
 
 const promptToAddStudioToken = async () => {
   const response = await Toast.askShowOrCloseOrNever(

extension/src/experiments/index.ts

@@ -641,6 +641,10 @@ export class Experiments extends BaseRepository<TableData> {
     this.notifyChanged()
   }
 
+  public transferExperimentDetails(oldName: string, newName: string) {
+    return this.experiments.transferDetails(oldName, newName)
+  }
+
   protected sendInitialWebviewData() {
     return this.webviewMessages.sendWebviewMessage()
   }

extension/src/experiments/model/index.ts

@@ -213,6 +213,20 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.coloredStatus[id]
   }
 
+  public transferDetails(oldName: string, newName: string) {
+    const starred = this.starredExperiments[oldName]
+    const selected = this.coloredStatus[oldName]
+
+    this.starredExperiments[newName] = starred
+    this.coloredStatus[newName] = selected
+
+    delete this.starredExperiments[oldName]
+    delete this.coloredStatus[oldName]
+
+    this.persistStatus()
+    this.persistStars()
+  }
+
   public hasRunningExperiment() {
     return this.running.length > 0
   }

extension/src/experiments/workspace.ts

@@ -241,7 +241,7 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
   }
 
   public async getInputAndRun(
-    runCommand: (...args: Args) => Promise<void> | void,
+    runCommand: (...args: Args) => Promise<void | string>,
     title: Title,
     defaultValue: string,
     ...args: Args
@@ -273,9 +273,10 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
   }
 
   public runCommand(commandId: CommandId, cwd: string, ...args: Args) {
-    return Toast.showOutput(
-      this.internalCommands.executeCommand(commandId, cwd, ...args)
-    )
+    const output = this.internalCommands.executeCommand(commandId, cwd, ...args)
+
+    void Toast.showOutput(output)
+    return output
   }
 
   public createRepository(

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

@@ -562,8 +562,13 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to rename an experiment', async () => {
-      const { mockMessageReceived } =
+      const { mockMessageReceived, experimentsModel } =
         await stubWorkspaceGettersWebview(disposable)
+      const transferDetailsSpy = spy(experimentsModel, 'transferDetails')
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const persistStatusSpy = spy(experimentsModel as any, 'persistStatus')
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      const persistStarsSpy = spy(experimentsModel as any, 'persistStars')
 
       const mockNewExperimentName = 'new-experiment-name'
       const inputEvent = getInputBoxEvent(mockNewExperimentName)
@@ -575,11 +580,15 @@ suite('Experiments Test Suite', () => {
       const mockRenameCalled = new Promise(resolve =>
         mockRenameExperiment.callsFake(() => {
           resolve(undefined)
-          return Promise.resolve('Renamed experiments:')
+          return Promise.resolve('Experiment renamed successfully')
         })
       )
 
-      const mockExperimentId = 'exp-e7a67'
+      const mockExperimentId = 'exp-83425'
+      const selectedColor = experimentsModel
+        .getSelectedRevisions()
+        .find(({ id }) => id === mockExperimentId)?.displayColor
+      expect(selectedColor).not.to.be.undefined
 
       mockMessageReceived.fire({
         payload: mockExperimentId,
@@ -593,6 +602,14 @@ suite('Experiments Test Suite', () => {
         mockExperimentId,
         mockNewExperimentName
       )
+      expect(transferDetailsSpy).to.be.calledOnce
+      expect(persistStatusSpy).to.be.calledOnce
+      expect(persistStarsSpy).to.be.calledOnce
+
+      expect(
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        (experimentsModel as any).coloredStatus[mockNewExperimentName]
+      ).to.equal(selectedColor)
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should handle a message to show the logs of an experiment', async () => {