Skip to content

Commit 94ffe7a

Browse files
seneyuseneyu
committed
Refactor tests with afterAll
1 parent 8719cde commit 94ffe7a

File tree

2 files changed

+45
-57
lines changed

2 files changed

+45
-57
lines changed

‎mlflow/tests/RunClient.test.ts‎

Lines changed: 31 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,11 @@
1-
import { describe, test, expect, beforeAll, beforeEach } from '@jest/globals';
1+
import {
2+
describe,
3+
test,
4+
expect,
5+
beforeAll,
6+
beforeEach,
7+
afterAll,
8+
} from '@jest/globals';
29
import RunClient from '../src/tracking/RunClient';
310
import ExperimentClient from '../src/tracking/ExperimentClient';
411
import {
@@ -13,6 +20,8 @@ describe('RunClient', () => {
1320
let runClient: RunClient;
1421
let experimentClient: ExperimentClient;
1522
let experimentId: string;
23+
let run: Run;
24+
const testIds: string[] = [];
1625

1726
beforeAll(async () => {
1827
await new Promise((resolve) => setTimeout(resolve, 2000));
@@ -24,12 +33,22 @@ describe('RunClient', () => {
2433
experimentId = await experimentClient.createExperiment(
2534
`Testing ${timestamp}`
2635
);
36+
testIds.push(experimentId);
37+
});
38+
39+
beforeEach(async () => {
40+
run = (await runClient.createRun(experimentId)) as Run;
41+
});
42+
43+
afterAll(async () => {
44+
for (const testId of testIds) {
45+
await experimentClient.deleteExperiment(testId);
46+
}
2747
});
2848

2949
// POST - Create a new run within an experiment
3050
describe('createRun', () => {
3151
test('- Should create a run with experiment_id', async () => {
32-
const run = (await runClient.createRun(experimentId)) as Run;
3352
expect(run.info.experiment_id).toBe(experimentId);
3453
});
3554

@@ -101,8 +120,6 @@ describe('RunClient', () => {
101120
// DELETE - Mark a run for deletion
102121
describe('deleteRun', () => {
103122
test('- Should delete a run with run_id', async () => {
104-
const run = (await runClient.createRun(experimentId)) as Run;
105-
106123
await expect(runClient.deleteRun(run.info.run_id)).resolves.not.toThrow();
107124

108125
// check if the run's lifecycle_stage has changed to "deleted"
@@ -128,11 +145,6 @@ describe('RunClient', () => {
128145

129146
// POST - Restore a deleted run
130147
describe('restoreRun', () => {
131-
let run: Run;
132-
beforeEach(async () => {
133-
run = (await runClient.createRun(experimentId)) as Run;
134-
});
135-
136148
test('- Should restore a deleted run with run_id', async () => {
137149
await runClient.deleteRun(run.info.run_id);
138150

@@ -180,8 +192,6 @@ describe('RunClient', () => {
180192
// GET - Get metadata, metrics, params, and tags for a run
181193
describe('getRun', () => {
182194
test('- Should retrieve metadata for a run with run_id', async () => {
183-
const run = (await runClient.createRun(experimentId)) as Run;
184-
185195
// create dummy data for created run
186196
const metrics: Metrics[] = [
187197
{ key: 'accuracy', value: 0.83, timestamp: 1694000700000 },
@@ -240,12 +250,6 @@ describe('RunClient', () => {
240250

241251
// POST - Update run metadata
242252
describe('updateRun', () => {
243-
let run: Run;
244-
245-
beforeEach(async () => {
246-
run = (await runClient.createRun(experimentId)) as Run;
247-
});
248-
249253
// parameterized testing for input status
250254
const allStatuses = [
251255
'RUNNING',
@@ -330,14 +334,9 @@ describe('RunClient', () => {
330334

331335
// POST - Log a metric for a run
332336
describe('logMetric', () => {
333-
let run: Run;
334337
const key = 'accuracy';
335338
const value = 0.9;
336339

337-
beforeEach(async () => {
338-
run = (await runClient.createRun(experimentId)) as Run;
339-
});
340-
341340
test('- Should log a metric with run_id, key, value, and timestamp', async () => {
342341
const timestamp = Date.now();
343342

@@ -405,12 +404,6 @@ describe('RunClient', () => {
405404

406405
// POST - Log a batch of metrics, params, and tags for a run
407406
describe('logBatch', () => {
408-
let run: Run;
409-
410-
beforeEach(async () => {
411-
run = (await runClient.createRun(experimentId)) as Run;
412-
});
413-
414407
test('- Should not throw error with just run_id', async () => {
415408
await expect(
416409
runClient.logBatch(run.info.run_id)
@@ -698,7 +691,6 @@ describe('RunClient', () => {
698691
];
699692

700693
test('- Should log inputs with run_id and datasets', async () => {
701-
const run = (await runClient.createRun(experimentId)) as Run;
702694
await runClient.logInputs(run.info.run_id, datasets);
703695

704696
// fetch run to confirm changes
@@ -709,8 +701,6 @@ describe('RunClient', () => {
709701
});
710702

711703
test('- Should handle errors and edge cases', async () => {
712-
const run = (await runClient.createRun(experimentId)) as Run;
713-
714704
// test with invalid_id
715705
const invalid_id = 'invalid_id';
716706
await expect(runClient.logInputs(invalid_id, datasets)).rejects.toThrow();
@@ -733,11 +723,10 @@ describe('RunClient', () => {
733723

734724
// POST - Set a tag on a run
735725
describe('setTag', () => {
736-
test('- Should set a tag on a run with run_id, key, and value', async () => {
737-
const run = (await runClient.createRun(experimentId)) as Run;
738-
const key = 'accuracy';
739-
const value = '0.99';
726+
const key = 'accuracy';
727+
const value = '0.99';
740728

729+
test('- Should set a tag on a run with run_id, key, and value', async () => {
741730
await runClient.setTag(run.info.run_id, key, value);
742731

743732
// fetch run to confirm changes
@@ -748,10 +737,6 @@ describe('RunClient', () => {
748737
});
749738

750739
test('- Should handle errors and edge cases', async () => {
751-
const run = (await runClient.createRun(experimentId)) as Run;
752-
const key = 'accuracy';
753-
const value = '0.99';
754-
755740
// test missing arguments
756741
// @ts-expect-error: testing for all missing arguments
757742
await expect(runClient.setTag()).rejects.toThrow();
@@ -793,8 +778,6 @@ describe('RunClient', () => {
793778
const key = 'test_key';
794779
const value = 'test_value';
795780
test('- Should delete a tag on a run with run_id and key', async () => {
796-
const run = (await runClient.createRun(experimentId)) as Run;
797-
798781
await runClient.setTag(run.info.run_id, key, value);
799782

800783
await runClient.deleteTag(run.info.run_id, key);
@@ -808,8 +791,6 @@ describe('RunClient', () => {
808791
});
809792

810793
test('- Should handle errors and edge cases', async () => {
811-
const run = (await runClient.createRun(experimentId)) as Run;
812-
813794
// testing missing arguments
814795
// @ts-expect-error: testing for all missing arguments
815796
await expect(runClient.deleteTag()).rejects.toThrow();
@@ -844,11 +825,10 @@ describe('RunClient', () => {
844825

845826
// POST - Log a param used for a run
846827
describe('logParam', () => {
847-
test('- Should log a param used for a run with run_id, key, and value', async () => {
848-
const run = (await runClient.createRun(experimentId)) as Run;
828+
const key = 'learning_rate';
829+
const value = '0.001';
849830

850-
const key = 'learning_rate';
851-
const value = '0.001';
831+
test('- Should log a param used for a run with run_id, key, and value', async () => {
852832
await runClient.logParam(run.info.run_id, key, value);
853833

854834
// fetch run to confirm changes
@@ -859,10 +839,6 @@ describe('RunClient', () => {
859839
});
860840

861841
test('- Should handle errors and edge cases', async () => {
862-
const run = (await runClient.createRun(experimentId)) as Run;
863-
const key = 'learning_rate';
864-
const value = '0.001';
865-
866842
// @ts-expect-error: testing for all missing arguments
867843
await expect(runClient.logParam()).rejects.toThrow();
868844
// @ts-expect-error: testing for missing key and value
@@ -886,11 +862,10 @@ describe('RunClient', () => {
886862

887863
// Get a list of all valuse for the specified metric for a given run
888864
describe('getMetricHisotry', () => {
889-
test('- Should get a list of all values for the specified metric for a given run with run_id and metric_key', async () => {
890-
const run = (await runClient.createRun(experimentId)) as Run;
891-
const key = 'accuracy';
892-
const value = 0.95;
865+
const key = 'accuracy';
866+
const value = 0.95;
893867

868+
test('- Should get a list of all values for the specified metric for a given run with run_id and metric_key', async () => {
894869
await runClient.logMetric(run.info.run_id, key, value);
895870
const metricHistory = (await runClient.getMetricHistory(
896871
run.info.run_id,
@@ -1008,6 +983,7 @@ describe('RunClient', () => {
1008983
const searchRunsExpId = await experimentClient.createExperiment(
1009984
`Search Runs Test ${Date.now()}`
1010985
);
986+
testIds.push(searchRunsExpId);
1011987

1012988
const runA = (await runClient.createRun(searchRunsExpId)) as Run;
1013989
await runClient.logMetric(runA.info.run_id, 'metric', 1.0);
@@ -1067,7 +1043,6 @@ describe('RunClient', () => {
10671043
// List artifacts for a run
10681044
describe('listArtifacts', () => {
10691045
test('- Should list artifacts with run_id', async () => {
1070-
const run = (await runClient.createRun(experimentId)) as Run;
10711046
const artifacts = await runClient.listArtifacts(run.info.run_id);
10721047

10731048
expect(artifacts).toHaveProperty('root_uri');

‎mlflow/tests/RunManager.test.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { describe, test, expect, beforeEach } from '@jest/globals';
1+
import { describe, test, expect, beforeEach, afterAll } from '@jest/globals';
22
import RunManager from '../src/workflows/RunManager';
33
import RunClient from '../src/tracking/RunClient';
44
import ExperimentClient from '../src/tracking/ExperimentClient';
@@ -10,6 +10,7 @@ describe('RunManager', () => {
1010
let experimentClient: ExperimentClient;
1111
let experimentId: string;
1212
const runIds: string[] = [];
13+
const experimentsToDelete: string[] = [];
1314

1415
beforeEach(async () => {
1516
await new Promise((resolve) => setTimeout(resolve, 2000));
@@ -18,12 +19,23 @@ describe('RunManager', () => {
1819
runManager = new RunManager('http://127.0.0.1:5002');
1920
});
2021

22+
afterAll(async () => {
23+
for (const runId of runIds) {
24+
await runClient.deleteRun(runId);
25+
}
26+
27+
for (const expId of experimentsToDelete) {
28+
await experimentClient.deleteExperiment(expId);
29+
}
30+
});
31+
2132
describe('cleanupRuns', () => {
2233
beforeEach(async () => {
2334
const timestamp = Date.now();
2435
experimentId = await experimentClient.createExperiment(
2536
`Testing ${timestamp}`
2637
);
38+
experimentsToDelete.push(experimentId);
2739

2840
// create test runs
2941
const createRun = async (metricKey: string, metricValue: number) => {
@@ -133,6 +145,7 @@ describe('RunManager', () => {
133145
targetExperimentId = await experimentClient.createExperiment(
134146
`Target Exp ${timestamp}`
135147
);
148+
experimentsToDelete.push(sourceExperimentId, targetExperimentId);
136149

137150
// log data for original run
138151
const run = (await runClient.createRun(sourceExperimentId)) as Run;

0 commit comments

Comments
 (0)