1
- import { describe , test , expect , beforeAll , jest } from '@jest/globals' ;
1
+ import { describe , test , expect , beforeAll , jest , afterAll } from '@jest/globals' ;
2
2
import RunClient from '../src/tracking/RunClient' ;
3
3
import ModelManager from '../src/workflows/ModelManager' ;
4
4
import ModelRegistryClient from '../src/model-registry/ModelRegistryClient' ;
@@ -29,7 +29,7 @@ describe('ModelManager', () => {
29
29
30
30
describe ( 'createRegisteredModelWithVersion' , ( ) => {
31
31
test ( 'Should create a new registered model with a version 1' , async ( ) => {
32
- const modelName2 = `createRegisteredModelWithVersion-test- model-${ timestamp } ` ;
32
+ const modelName2 = `createRegisteredModelWithVersion-model-${ timestamp } ` ;
33
33
const newModelVersion : keyable =
34
34
await modelManager . createRegisteredModelWithVersion (
35
35
modelName2 ,
@@ -38,13 +38,14 @@ describe('ModelManager', () => {
38
38
) ;
39
39
expect ( newModelVersion . name ) . toBe ( modelName2 ) ;
40
40
expect ( newModelVersion . version ) . toBe ( '1' ) ;
41
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
41
42
} ) ;
42
43
test ( 'Should throw an error for duplicate model name' , async ( ) => {
43
44
// Making it so console.error doesn't show up in the console
44
45
const consoleErrorMock = jest
45
46
. spyOn ( global . console , 'error' )
46
47
. mockImplementation ( ( ) => { } ) ;
47
- const modelName3 = `createRegisteredModelWithVersionDup-test- model-${ timestamp } ` ;
48
+ const modelName3 = `createRegisteredModelWithVersionDup-model-${ timestamp } ` ;
48
49
await modelManager . createRegisteredModelWithVersion (
49
50
modelName3 ,
50
51
run . info . artifact_uri ,
@@ -62,6 +63,7 @@ describe('ModelManager', () => {
62
63
expect ( consoleErrorMock ) . toHaveBeenCalled ( ) ;
63
64
// restoring the console to show console.error again to not affect other tests
64
65
consoleErrorMock . mockRestore ( ) ;
66
+ await modelRegistryClient . deleteRegisteredModel ( modelName3 ) ;
65
67
} ) ;
66
68
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
67
69
// Making it so console.error doesn't show up in the console
@@ -81,7 +83,7 @@ describe('ModelManager', () => {
81
83
82
84
describe ( 'updateRegisteredModelDescriptionAndTag' , ( ) => {
83
85
test ( "Should update a registered model's description and tag" , async ( ) => {
84
- const modelName2 = `updateRegisteredModelDescriptionAndTag-test- model-${ timestamp } ` ;
86
+ const modelName2 = `updateRegisteredModelDescriptionAndTag-model-${ timestamp } ` ;
85
87
const updateRegisteredModelDescriptionAndTagObject : keyable =
86
88
await modelRegistryClient . createRegisteredModel ( modelName2 ) ;
87
89
expect ( updateRegisteredModelDescriptionAndTagObject . description ) . toBe (
@@ -99,6 +101,7 @@ describe('ModelManager', () => {
99
101
) ;
100
102
expect ( updatedModelDescriptionAndTag . tags [ 0 ] . key ) . toBe ( 'modelTagKey' ) ;
101
103
expect ( updatedModelDescriptionAndTag . tags [ 0 ] . value ) . toBe ( 'modelTagValue' ) ;
104
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
102
105
} ) ;
103
106
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
104
107
// Making it so console.error doesn't show up in the console
@@ -118,7 +121,7 @@ describe('ModelManager', () => {
118
121
119
122
describe ( 'updateAllLatestModelVersion' , ( ) => {
120
123
test ( "Should update the latest model version's alias, tag, and description" , async ( ) => {
121
- const modelName2 = `updateAllLatestModelVersion-test- model-${ timestamp } ` ;
124
+ const modelName2 = `updateAllLatestModelVersion-model-${ timestamp } ` ;
122
125
await modelManager . createRegisteredModelWithVersion (
123
126
modelName2 ,
124
127
run . info . artifact_uri ,
@@ -151,6 +154,7 @@ describe('ModelManager', () => {
151
154
expect ( updateAllLatestModelVersionObject . tags [ 0 ] . value ) . toBe (
152
155
'modelVersionTagValue'
153
156
) ;
157
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
154
158
} ) ;
155
159
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
156
160
// Making it so console.error doesn't show up in the console
@@ -170,7 +174,7 @@ describe('ModelManager', () => {
170
174
171
175
describe ( 'setLatestModelVersionTag' , ( ) => {
172
176
test ( 'Should add a new tag key/value for the latest version of the specified registered model' , async ( ) => {
173
- const modelName2 = `setLatestModelVersionTag-test- model-${ timestamp } ` ;
177
+ const modelName2 = `setLatestModelVersionTag-model-${ timestamp } ` ;
174
178
await modelManager . createRegisteredModelWithVersion (
175
179
modelName2 ,
176
180
run . info . artifact_uri ,
@@ -192,6 +196,7 @@ describe('ModelManager', () => {
192
196
expect ( latestModelVersion [ 0 ] . name ) . toBe ( modelName2 ) ;
193
197
expect ( latestModelVersion [ 0 ] . tags [ 0 ] . key ) . toBe ( 'modelVersionTagKey' ) ;
194
198
expect ( latestModelVersion [ 0 ] . tags [ 0 ] . value ) . toBe ( 'modelVersionTagValue' ) ;
199
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
195
200
} ) ;
196
201
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
197
202
// Making it so console.error doesn't show up in the console
@@ -211,7 +216,7 @@ describe('ModelManager', () => {
211
216
212
217
describe ( 'setLatestModelVersionAlias' , ( ) => {
213
218
test ( 'Should add an alias for the latest version of the specified registered model' , async ( ) => {
214
- const modelName2 = `setLatestModelVersionAlias-test- model-${ timestamp } ` ;
219
+ const modelName2 = `setLatestModelVersionAlias-model-${ timestamp } ` ;
215
220
await modelManager . createRegisteredModelWithVersion (
216
221
modelName2 ,
217
222
run . info . artifact_uri ,
@@ -231,6 +236,7 @@ describe('ModelManager', () => {
231
236
expect ( latestModelVersion [ 0 ] . version ) . toBe ( '2' ) ;
232
237
expect ( latestModelVersion [ 0 ] . name ) . toBe ( modelName2 ) ;
233
238
expect ( latestModelVersion [ 0 ] . aliases [ 0 ] ) . toBe ( 'modelVersionAlias' ) ;
239
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
234
240
} ) ;
235
241
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
236
242
// Making it so console.error doesn't show up in the console
@@ -250,7 +256,7 @@ describe('ModelManager', () => {
250
256
251
257
describe ( 'updateLatestModelVersion' , ( ) => {
252
258
test ( 'Should update the description of the latest model version' , async ( ) => {
253
- const modelName2 = `updateLatestModelVersion-test- model-${ timestamp } ` ;
259
+ const modelName2 = `updateLatestModelVersion-model-${ timestamp } ` ;
254
260
await modelManager . createRegisteredModelWithVersion (
255
261
modelName2 ,
256
262
run . info . artifact_uri ,
@@ -269,6 +275,7 @@ describe('ModelManager', () => {
269
275
expect ( latestModelVersion . version ) . toBe ( '2' ) ;
270
276
expect ( latestModelVersion . name ) . toBe ( modelName2 ) ;
271
277
expect ( latestModelVersion . description ) . toBe ( 'modelDescription' ) ;
278
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
272
279
} ) ;
273
280
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
274
281
// Making it so console.error doesn't show up in the console
@@ -288,7 +295,7 @@ describe('ModelManager', () => {
288
295
289
296
describe ( 'updateAllModelVersion' , ( ) => {
290
297
test ( "Should update the specified version's alias, tag, and description" , async ( ) => {
291
- const modelName2 = `updateAllModelVersion-test- model-${ timestamp } ` ;
298
+ const modelName2 = `updateAllModelVersion-model-${ timestamp } ` ;
292
299
await modelManager . createRegisteredModelWithVersion (
293
300
modelName2 ,
294
301
run . info . artifact_uri ,
@@ -315,6 +322,7 @@ describe('ModelManager', () => {
315
322
expect ( updatedModelVersionAll . description ) . toBe (
316
323
'modelVersionDescription'
317
324
) ;
325
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
318
326
} ) ;
319
327
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
320
328
// Making it so console.error doesn't show up in the console
@@ -334,7 +342,7 @@ describe('ModelManager', () => {
334
342
335
343
describe ( 'deleteLatestModelVersion' , ( ) => {
336
344
test ( 'Should delete the latest version of the model' , async ( ) => {
337
- const modelName2 = `deleteLatestModelVersion-test- model-${ timestamp } ` ;
345
+ const modelName2 = `deleteLatestModelVersion-model-${ timestamp } ` ;
338
346
await modelManager . createRegisteredModelWithVersion (
339
347
modelName2 ,
340
348
run . info . artifact_uri ,
@@ -349,6 +357,7 @@ describe('ModelManager', () => {
349
357
const latestModelVersion : keyable =
350
358
await modelRegistryClient . getLatestModelVersions ( modelName2 ) ;
351
359
expect ( latestModelVersion [ 0 ] . version ) . toBe ( '1' ) ;
360
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
352
361
} ) ;
353
362
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
354
363
// Making it so console.error doesn't show up in the console
@@ -374,7 +383,7 @@ describe('ModelManager', () => {
374
383
375
384
const runData : keyable = await runClient . getRun ( run . info . run_id ) ;
376
385
377
- const modelName2 = `createModelFromRunWithBestMetric-test- modelMax-${ timestamp } ` ;
386
+ const modelName2 = `createModelFromRunWithBestMetric-modelMax-${ timestamp } ` ;
378
387
await modelManager . createModelFromRunWithBestMetric (
379
388
[ run . info . experiment_id ] ,
380
389
runData . data . metrics [ 0 ] . key ,
@@ -387,6 +396,8 @@ describe('ModelManager', () => {
387
396
expect ( maxModel . latest_versions [ 0 ] . source ) . toBe (
388
397
runData . info . artifact_uri
389
398
) ;
399
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
400
+ await runClient . deleteRun ( run2 . info . run_id ) ;
390
401
} ) ;
391
402
392
403
test ( 'Should create model from run with best(minimum) metric' , async ( ) => {
@@ -396,7 +407,7 @@ describe('ModelManager', () => {
396
407
397
408
const runData2 : keyable = await runClient . getRun ( run2 . info . run_id ) ;
398
409
399
- const modelName2 = `createModelFromRunWithBestMetric-test- modelMin-${ timestamp } ` ;
410
+ const modelName2 = `createModelFromRunWithBestMetric-modelMin-${ timestamp } ` ;
400
411
await modelManager . createModelFromRunWithBestMetric (
401
412
[ run2 . info . experiment_id ] ,
402
413
runData2 . data . metrics [ 0 ] . key ,
@@ -409,6 +420,8 @@ describe('ModelManager', () => {
409
420
expect ( minModel . latest_versions [ 0 ] . source ) . toBe (
410
421
runData2 . info . artifact_uri
411
422
) ;
423
+ await modelRegistryClient . deleteRegisteredModel ( modelName2 ) ;
424
+ await runClient . deleteRun ( run2 . info . run_id ) ;
412
425
} ) ;
413
426
414
427
test ( 'Should throw and console.error if no parameters are passed to the method' , async ( ) => {
@@ -426,4 +439,8 @@ describe('ModelManager', () => {
426
439
consoleErrorMock . mockRestore ( ) ;
427
440
} ) ;
428
441
} ) ;
442
+
443
+ afterAll ( async ( ) => {
444
+ await runClient . deleteRun ( run . info . run_id ) ;
445
+ } ) ;
429
446
} ) ;
0 commit comments