Skip to content

Commit 379566e

Browse files
sroy3sroy3
authored
Check stages for all run/queue commands (#3304)
* Chack for pipeline stages for all run/queue actions * Fix some tests * Fix other test * Add tests * Apply review comments * Revert "Apply review comments" This reverts commit 3382d04.
1 parent aa587d7 commit 379566e

File tree

7 files changed

+148
-62
lines changed

7 files changed

+148
-62
lines changed

‎extension/src/experiments/commands/register.ts‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -289,7 +289,7 @@ const registerExperimentRunCommands = (
289289
): void => {
290290
internalCommands.registerExternalCliCommand(
291291
RegisteredCliCommands.EXPERIMENT_RUN,
292-
() => experiments.getCwdThenRun(AvailableCommands.EXPERIMENT_RUN, true)
292+
() => experiments.getCwdThenRun(AvailableCommands.EXPERIMENT_RUN)
293293
)
294294

295295
internalCommands.registerExternalCliCommand(
@@ -299,11 +299,7 @@ const registerExperimentRunCommands = (
299299

300300
internalCommands.registerExternalCliCommand(
301301
RegisteredCliCommands.EXPERIMENT_RESET_AND_RUN,
302-
() =>
303-
experiments.getCwdThenRun(
304-
AvailableCommands.EXPERIMENT_RESET_AND_RUN,
305-
true
306-
)
302+
() => experiments.getCwdThenRun(AvailableCommands.EXPERIMENT_RESET_AND_RUN)
307303
)
308304

309305
internalCommands.registerExternalCliCommand(

‎extension/src/experiments/workspace.test.ts‎

Lines changed: 75 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ describe('Experiments', () => {
8989
describe('getCwdThenReport', () => {
9090
it('should call the correct function with the correct parameters if a project is picked', async () => {
9191
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
92+
mockedListStages.mockResolvedValueOnce('train')
9293

9394
await workspaceExperiments.getCwdThenReport(mockedCommandId)
9495

@@ -105,6 +106,15 @@ describe('Experiments', () => {
105106
expect(mockedQuickPickOne).toHaveBeenCalledTimes(1)
106107
expect(mockedExpFunc).not.toHaveBeenCalled()
107108
})
109+
110+
it('should check and ask for the creation of a pipeline stage before running the command', async () => {
111+
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
112+
mockedListStages.mockResolvedValueOnce('')
113+
114+
await workspaceExperiments.getCwdThenReport(mockedCommandId)
115+
116+
expect(mockedExpFunc).not.toHaveBeenCalled()
117+
})
108118
})
109119

110120
describe('pauseUpdatesThenRun', () => {
@@ -123,6 +133,7 @@ describe('Experiments', () => {
123133
describe('getExpNameThenRun', () => {
124134
it('should call the correct function with the correct parameters if a project and experiment are picked', async () => {
125135
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
136+
mockedListStages.mockResolvedValueOnce('train')
126137
mockedPickExperiment.mockResolvedValueOnce({
127138
id: 'a123456',
128139
name: 'exp-123'
@@ -144,11 +155,21 @@ describe('Experiments', () => {
144155
expect(mockedQuickPickOne).toHaveBeenCalledTimes(1)
145156
expect(mockedExpFunc).not.toHaveBeenCalled()
146157
})
158+
159+
it('should check and ask for the creation of a pipeline stage before running the command', async () => {
160+
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
161+
mockedListStages.mockResolvedValueOnce('')
162+
163+
await workspaceExperiments.getCwdAndExpNameThenRun(mockedCommandId)
164+
165+
expect(mockedExpFunc).not.toHaveBeenCalled()
166+
})
147167
})
148168

149169
describe('getCwdAndQuickPickThenRun', () => {
150170
it('should call the correct function with the correct parameters if a project and experiment are picked and the quick pick returns a list', async () => {
151171
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
172+
mockedListStages.mockResolvedValueOnce('train')
152173

153174
const mockedPickedOptions = ['a', 'b', 'c']
154175
const mockedQuickPick = jest
@@ -184,6 +205,7 @@ describe('Experiments', () => {
184205

185206
it('should not call the function if quick picks are not provided', async () => {
186207
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
208+
mockedListStages.mockResolvedValueOnce('train')
187209
const mockedQuickPick = jest.fn().mockResolvedValueOnce(undefined)
188210

189211
await workspaceExperiments.getCwdAndQuickPickThenRun(
@@ -195,11 +217,28 @@ describe('Experiments', () => {
195217
expect(mockedQuickPick).toHaveBeenCalledTimes(1)
196218
expect(mockedExpFunc).not.toHaveBeenCalled()
197219
})
220+
221+
it('should check and ask for the creation of a pipeline stage before running the command', async () => {
222+
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
223+
mockedListStages.mockResolvedValueOnce('')
224+
const mockedPickedOptions = ['a', 'b', 'c']
225+
const mockedQuickPick = jest
226+
.fn()
227+
.mockResolvedValueOnce(mockedPickedOptions)
228+
229+
await workspaceExperiments.getCwdAndQuickPickThenRun(
230+
mockedCommandId,
231+
mockedQuickPick
232+
)
233+
234+
expect(mockedExpFunc).not.toHaveBeenCalled()
235+
})
198236
})
199237

200238
describe('getCwdExpNameAndInputThenRun', () => {
201239
it('should call the correct function with the correct parameters if a project and experiment are picked and an input provided', async () => {
202240
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
241+
mockedListStages.mockResolvedValueOnce('train')
203242
mockedPickExperiment.mockResolvedValueOnce({
204243
id: 'a123456',
205244
name: 'exp-123'
@@ -238,6 +277,7 @@ describe('Experiments', () => {
238277

239278
it('should not call the function if user input is not provided', async () => {
240279
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
280+
mockedListStages.mockResolvedValueOnce('train')
241281
mockedPickExperiment.mockResolvedValueOnce({
242282
id: 'b456789',
243283
name: 'exp-456'
@@ -254,11 +294,30 @@ describe('Experiments', () => {
254294
expect(mockedGetInput).toHaveBeenCalledTimes(1)
255295
expect(mockedExpFunc).not.toHaveBeenCalled()
256296
})
297+
298+
it('should check and ask for the creation of a pipeline stage before running the command', async () => {
299+
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
300+
mockedListStages.mockResolvedValueOnce('')
301+
mockedPickExperiment.mockResolvedValueOnce({
302+
id: 'a123456',
303+
name: 'exp-123'
304+
})
305+
mockedGetInput.mockResolvedValueOnce('abc123')
306+
307+
await workspaceExperiments.getCwdExpNameAndInputThenRun(
308+
(cwd: string, ...args: Args) =>
309+
workspaceExperiments.runCommand(mockedCommandId, cwd, ...args),
310+
'enter your password please' as Title
311+
)
312+
313+
expect(mockedExpFunc).not.toHaveBeenCalled()
314+
})
257315
})
258316

259317
describe('getCwdThenRun', () => {
260318
it('should call the correct function with the correct parameters if a project is picked', async () => {
261319
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
320+
mockedListStages.mockResolvedValueOnce('train')
262321

263322
await workspaceExperiments.getCwdThenRun(mockedCommandId)
264323

@@ -276,25 +335,17 @@ describe('Experiments', () => {
276335
expect(mockedExpFunc).not.toHaveBeenCalled()
277336
})
278337

279-
it('should ensure that a dvc.yaml file exists if the registered command needs it', async () => {
338+
it('should ensure that a dvc.yaml file exists', async () => {
280339
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
281340
mockedGetValidInput.mockResolvedValueOnce('train')
282341
mockedListStages.mockResolvedValueOnce('')
283342
mockedQuickPickOneOrInput.mockResolvedValueOnce(
284343
'path/to/training_script.py'
285344
)
286345

287-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
288-
289-
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledTimes(1)
290-
})
291-
292-
it('should not ensure that a dvc.yaml file exists if the registered command does not require it', async () => {
293-
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
294-
295346
await workspaceExperiments.getCwdThenRun(mockedCommandId)
296347

297-
expect(mockedFindOrCreateDvcYamlFile).not.toHaveBeenCalled()
348+
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledTimes(1)
298349
})
299350

300351
it('should check for pipelines when a command needs it and continue with the command if there is a pipeline', async () => {
@@ -307,7 +358,7 @@ describe('Experiments', () => {
307358
mockedListStages.mockResolvedValueOnce('train')
308359
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
309360

310-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
361+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
311362

312363
expect(executeCommandSpy).toHaveBeenCalledWith(
313364
AvailableCommands.STAGE_LIST,
@@ -323,7 +374,7 @@ describe('Experiments', () => {
323374
mockedListStages.mockResolvedValueOnce('')
324375
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
325376

326-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
377+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
327378

328379
expect(mockedGetValidInput).toHaveBeenCalledWith(
329380
Title.ENTER_STAGE_NAME,
@@ -336,7 +387,7 @@ describe('Experiments', () => {
336387
mockedListStages.mockResolvedValueOnce('train')
337388
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
338389

339-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
390+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
340391

341392
expect(mockedGetValidInput).not.toHaveBeenCalledWith(
342393
Title.ENTER_STAGE_NAME,
@@ -355,7 +406,7 @@ describe('Experiments', () => {
355406
mockedListStages.mockResolvedValueOnce('')
356407
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
357408

358-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
409+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
359410

360411
expect(executeCommandSpy).not.toHaveBeenCalledWith(
361412
mockedCommandId,
@@ -371,7 +422,7 @@ describe('Experiments', () => {
371422
'path/to/training_script.py'
372423
)
373424

374-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
425+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
375426

376427
expect(mockedQuickPickOneOrInput).toHaveBeenCalledTimes(1)
377428
})
@@ -385,7 +436,7 @@ describe('Experiments', () => {
385436
mockedQuickPickOneOrInput.mockResolvedValueOnce(trainingScript)
386437
mockedGetFileExtension.mockReturnValueOnce('.py')
387438

388-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
439+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
389440

390441
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith(
391442
mockedDvcRoot,
@@ -408,7 +459,7 @@ describe('Experiments', () => {
408459
'path/to/training_script.py'
409460
)
410461

411-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
462+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
412463

413464
expect(executeCommandSpy).toHaveBeenCalledWith(
414465
mockedCommandId,
@@ -425,7 +476,7 @@ describe('Experiments', () => {
425476
)
426477
mockedGetFileExtension.mockReturnValueOnce('.py')
427478

428-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
479+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
429480

430481
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith(
431482
mockedDvcRoot,
@@ -444,7 +495,7 @@ describe('Experiments', () => {
444495
)
445496
mockedGetFileExtension.mockReturnValueOnce('.ipynb')
446497

447-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
498+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
448499

449500
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith(
450501
mockedDvcRoot,
@@ -463,7 +514,7 @@ describe('Experiments', () => {
463514
)
464515
mockedGetFileExtension.mockReturnValueOnce('.ipynb')
465516

466-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
517+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
467518

468519
expect(mockedGetInput).not.toHaveBeenCalledWith(
469520
Title.ENTER_COMMAND_TO_RUN
@@ -479,7 +530,7 @@ describe('Experiments', () => {
479530
)
480531
mockedGetFileExtension.mockReturnValueOnce('.js')
481532

482-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
533+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
483534

484535
expect(mockedGetInput).toHaveBeenCalledWith(Title.ENTER_COMMAND_TO_RUN)
485536
})
@@ -496,7 +547,7 @@ describe('Experiments', () => {
496547
mockedGetFileExtension.mockReturnValueOnce('.js')
497548
mockedGetInput.mockResolvedValueOnce(customCommand)
498549

499-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
550+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
500551

501552
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith(
502553
mockedDvcRoot,
@@ -516,7 +567,7 @@ describe('Experiments', () => {
516567
mockedGetFileExtension.mockReturnValueOnce('.js')
517568
mockedGetInput.mockResolvedValueOnce(undefined)
518569

519-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
570+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
520571

521572
expect(mockedFindOrCreateDvcYamlFile).toHaveBeenCalledWith(
522573
mockedDvcRoot,
@@ -537,7 +588,7 @@ describe('Experiments', () => {
537588
mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
538589
mockedQuickPickOneOrInput.mockResolvedValueOnce('')
539590

540-
await workspaceExperiments.getCwdThenRun(mockedCommandId, true)
591+
await workspaceExperiments.getCwdThenRun(mockedCommandId)
541592

542593
expect(executeCommandSpy).not.toHaveBeenCalledWith(
543594
mockedCommandId,

0 commit comments

Comments
 (0)