Skip to content

Commit 1d0fbf3

Browse files
austinbfraseraustinbfraser
committed
refactored ExperimentManager tests to use testUtils, added runProperties to testUtils, refactored ExperimentManager to not console error, changed order of imports in ExperimentClient test
1 parent 2f0c2c9 commit 1d0fbf3

File tree

4 files changed

+48
-60
lines changed

4 files changed

+48
-60
lines changed

‎mlflow/src/workflows/ExperimentManager.ts‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import RunClient from '@tracking/RunClient';
33
import { ApiError } from '@utils/apiError';
44

55
interface keyable {
6-
[key: string]: any
6+
[key: string]: any;
77
}
88

99
class ExperimentManager {
@@ -50,7 +50,10 @@ class ExperimentManager {
5050
): Promise<object> {
5151
try {
5252
// create run
53-
const run:keyable = await this.runClient.createRun(experiment_id, run_name);
53+
const run: keyable = await this.runClient.createRun(
54+
experiment_id,
55+
run_name
56+
);
5457
const run_id = run.info.run_id;
5558

5659
// log metric, params, and tags via logBatch
@@ -70,10 +73,8 @@ class ExperimentManager {
7073
return (latestRun as { run_info: object }).run_info;
7174
} catch (error) {
7275
if (error instanceof ApiError) {
73-
console.error(`API Error (${error.statusCode}): ${error.message}`);
7476
throw error;
7577
} else {
76-
console.error('An unexpected error occurred:', error);
7778
throw new Error();
7879
}
7980
}
@@ -118,7 +119,10 @@ class ExperimentManager {
118119
);
119120

120121
// create run
121-
const run:keyable = await this.runClient.createRun(experiment_id, run_name);
122+
const run: keyable = await this.runClient.createRun(
123+
experiment_id,
124+
run_name
125+
);
122126
const run_id = run.info.run_id;
123127

124128
// log metric, params, and tags via logBatch
@@ -138,10 +142,8 @@ class ExperimentManager {
138142
return (latestRun as { run_info: object }).run_info;
139143
} catch (error) {
140144
if (error instanceof ApiError) {
141-
console.error(`API Error (${error.statusCode}): ${error.message}`);
142145
throw error;
143146
} else {
144-
console.error('An unexpected error occurred:', error);
145147
throw new Error();
146148
}
147149
}
@@ -207,7 +209,7 @@ class ExperimentManager {
207209
if (order === 1 || order === 'DESC') orderString = 'DESC';
208210
else if (order === -1 || order === 'ASC') orderString = 'ASC';
209211
const arg = `metrics.${primaryMetric} ${orderString}`;
210-
const data:keyable = await this.runClient.searchRuns(
212+
const data: keyable = await this.runClient.searchRuns(
211213
[experiment_id],
212214
'',
213215
undefined,
@@ -234,10 +236,8 @@ class ExperimentManager {
234236
return data.runs;
235237
} catch (error) {
236238
if (error instanceof ApiError) {
237-
console.error(`API Error (${error.statusCode}): ${error.message}`);
238239
throw error;
239240
} else {
240-
console.error('An unexpected error occurred:', error);
241241
throw new Error();
242242
}
243243
}

‎mlflow/tests/ExperimentClient.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@ import ExperimentClient from '../src/tracking/ExperimentClient';
33
import { ApiError } from '../src/utils/apiError';
44
import {
55
createTestExperiment,
6-
experimentProperties,
76
deleteTestExperiments,
7+
experimentProperties,
88
ExpSearchResults,
99
TRACKING_SERVER_URI,
1010
} from './testUtils';

‎mlflow/tests/ExperimentManager.test.ts‎

Lines changed: 21 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,12 @@ import { describe, test, expect, beforeAll, afterAll } from '@jest/globals';
22
import { ApiError } from '../src/utils/apiError';
33
import ExperimentClient from '../src/tracking/ExperimentClient';
44
import ExperimentManager from '../src/workflows/ExperimentManager';
5+
import {
6+
createTestExperiment,
7+
deleteTestExperiments,
8+
runProperties,
9+
TRACKING_SERVER_URI,
10+
} from './testUtils';
511

612
describe('ExperimentManager', () => {
713
let experimentClient: ExperimentClient;
@@ -45,17 +51,14 @@ describe('ExperimentManager', () => {
4551
beforeAll(async () => {
4652
// Add a small delay to ensure MLflow is fully ready
4753
await new Promise((resolve) => setTimeout(resolve, 2000));
48-
experimentClient = new ExperimentClient('http://127.0.0.1:5002');
49-
experimentManager = new ExperimentManager('http://127.0.0.1:5002');
54+
experimentClient = new ExperimentClient(TRACKING_SERVER_URI);
55+
experimentManager = new ExperimentManager(TRACKING_SERVER_URI);
5056
});
5157

5258
describe('runExistingExperiment', () => {
5359
test('should run an existing experiment and return the run object', async () => {
54-
const num = Math.random().toString().slice(2, 11);
55-
const name = `Test experiment ${num}`;
56-
const exp = await experimentClient.createExperiment(name);
60+
const exp: string = await createTestExperiment();
5761
testIds.push(exp);
58-
5962
const run = await experimentManager.runExistingExperiment(
6063
exp,
6164
undefined,
@@ -64,24 +67,14 @@ describe('ExperimentManager', () => {
6467
tags,
6568
model
6669
);
67-
68-
expect(run).toHaveProperty('run_id');
69-
expect(run).toHaveProperty('run_uuid');
70-
expect(run).toHaveProperty('run_name');
71-
expect(run).toHaveProperty('experiment_id');
72-
expect(run).toHaveProperty('user_id');
73-
expect(run).toHaveProperty('status');
74-
expect(run).toHaveProperty('start_time');
75-
expect(run).toHaveProperty('artifact_uri');
76-
expect(run).toHaveProperty('lifecycle_stage');
70+
for (const property of runProperties) {
71+
expect(run).toHaveProperty(property);
72+
}
7773
});
7874

7975
test('should throw error if an invalid experiment ID is passed in', async () => {
80-
const num = Math.random().toString().slice(2, 11);
81-
const name = `Test experiment ${num}`;
82-
const exp = await experimentClient.createExperiment(name);
76+
const exp: string = await createTestExperiment();
8377
testIds.push(exp);
84-
8578
await expect(
8679
experimentManager.runExistingExperiment(
8780
'invalidExperimentId',
@@ -97,8 +90,7 @@ describe('ExperimentManager', () => {
9790

9891
describe('runNewExperiment', () => {
9992
test('should run a new experiment and return the run object', async () => {
100-
const num = Math.random().toString().slice(2, 11);
101-
const name = `Test experiment ${num}`;
93+
const name: string = `Test experiment ${Date.now()}`;
10294
const run: {
10395
experiment_id?: string;
10496
} = await experimentManager.runNewExperiment(
@@ -112,22 +104,14 @@ describe('ExperimentManager', () => {
112104
if (run.experiment_id) {
113105
testIds.push(run.experiment_id);
114106
}
115-
116-
expect(run).toHaveProperty('run_id');
117-
expect(run).toHaveProperty('run_uuid');
118-
expect(run).toHaveProperty('run_name');
119-
expect(run).toHaveProperty('experiment_id');
120-
expect(run).toHaveProperty('user_id');
121-
expect(run).toHaveProperty('status');
122-
expect(run).toHaveProperty('start_time');
123-
expect(run).toHaveProperty('artifact_uri');
124-
expect(run).toHaveProperty('lifecycle_stage');
107+
for (const property of runProperties) {
108+
expect(run).toHaveProperty(property);
109+
}
125110
});
126111

127112
test('should throw error if an invalid experiment name is passed in', async () => {
128-
const num = Math.random().toString().slice(2, 11);
129-
const name = `Test experiment ${num}`;
130-
const exp = await experimentClient.createExperiment(name);
113+
const name: string = `Test experiment ${Date.now()}`;
114+
const exp: string = await experimentClient.createExperiment(name);
131115
testIds.push(exp);
132116
await expect(
133117
experimentManager.runNewExperiment(
@@ -144,9 +128,7 @@ describe('ExperimentManager', () => {
144128

145129
describe('experimentSummary', () => {
146130
test("should return an array of all the passed-in experiment's runs, sorted according to the passed-in metric", async () => {
147-
const num = Math.random().toString().slice(2, 11);
148-
const name = `Test experiment ${num}`;
149-
const exp = await experimentClient.createExperiment(name);
131+
const exp: string = await createTestExperiment();
150132
testIds.push(exp);
151133
for (const metric of metricsAll) {
152134
await experimentManager.runExistingExperiment(
@@ -176,12 +158,5 @@ describe('ExperimentManager', () => {
176158
});
177159
});
178160

179-
afterAll(async () => {
180-
while (testIds.length > 0) {
181-
const id = testIds.pop();
182-
if (id) {
183-
await experimentClient.deleteExperiment(id);
184-
}
185-
}
186-
});
161+
afterAll(async () => await deleteTestExperiments(testIds));
187162
});

‎mlflow/tests/testUtils.ts‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,18 @@ export const experimentProperties: string[] = [
88
'artifact_location',
99
'lifecycle_stage',
1010
'last_update_time',
11-
'creation_time'
11+
'creation_time',
12+
];
13+
export const runProperties: string[] = [
14+
'run_id',
15+
'run_uuid',
16+
'run_name',
17+
'experiment_id',
18+
'user_id',
19+
'status',
20+
'start_time',
21+
'artifact_uri',
22+
'lifecycle_stage',
1223
];
1324
export type ExpSearchResults = {
1425
experiments?: Experiment[];
@@ -17,7 +28,9 @@ export type ExpSearchResults = {
1728

1829
const experimentClient = new ExperimentClient(TRACKING_SERVER_URI);
1930

20-
export const createTestExperiment = async (prefix = 'Test experiment'): Promise<string> => {
31+
export const createTestExperiment = async (
32+
prefix = 'Test experiment'
33+
): Promise<string> => {
2134
const timestamp = Date.now();
2235
return await experimentClient.createExperiment(`${prefix} ${timestamp}`);
2336
};
@@ -30,4 +43,4 @@ export const deleteTestExperiments = async (experimentIds: string[]) => {
3043
console.warn(`Failed to delete experiment ${id}: ${err}`);
3144
}
3245
}
33-
};
46+
};

0 commit comments

Comments
 (0)