Skip to content

Commit 3104110

Browse files
seneyuseneyu
committed
Merge branch 'testUtils' into stephany/jest
2 parents 94ffe7a + 1d0fbf3 commit 3104110

File tree

5 files changed

+148
-153
lines changed

5 files changed

+148
-153
lines changed

‎mlflow/src/tracking/ExperimentClient.ts

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,6 @@ class ExperimentClient {
180180
response.status
181181
);
182182
};
183-
184-
console.log(`Experiment ID ${experiment_id} successfully deleted`);
185183
}
186184

187185
/**
@@ -212,8 +210,6 @@ class ExperimentClient {
212210
response.status
213211
);
214212
}
215-
216-
console.log(`Experiment ID ${experiment_id} successfully restored`);
217213
}
218214

219215
/**
@@ -246,10 +242,6 @@ class ExperimentClient {
246242
response.status
247243
);
248244
}
249-
250-
console.log(
251-
`Experiment ID ${experiment_id} successfully updated - new name is ${new_name}`
252-
);
253245
}
254246

255247
/**
@@ -284,8 +276,6 @@ class ExperimentClient {
284276
response.status
285277
);
286278
}
287-
288-
console.log(`Set tag to experiment ID ${experiment_id} successfully`);
289279
}
290280
}
291281

‎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: 71 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1,32 +1,27 @@
11
import { describe, test, expect, beforeAll, afterAll } from '@jest/globals';
22
import ExperimentClient from '../src/tracking/ExperimentClient';
33
import { ApiError } from '../src/utils/apiError';
4-
import { Experiment } from '../src/utils/interface';
4+
import {
5+
createTestExperiment,
6+
deleteTestExperiments,
7+
experimentProperties,
8+
ExpSearchResults,
9+
TRACKING_SERVER_URI,
10+
} from './testUtils';
511

612
describe('ExperimentClient', () => {
713
let experimentClient: ExperimentClient;
8-
let experimentId: string;
9-
let experimentName: string;
1014
const testIds: string[] = [];
1115

1216
beforeAll(async () => {
1317
// Add a small delay to ensure MLflow is fully ready
1418
await new Promise((resolve) => setTimeout(resolve, 2000));
15-
experimentClient = new ExperimentClient('http://127.0.0.1:5002');
16-
17-
// Generate the experiment ID to be used generically in later tests
18-
const timestamp = Date.now();
19-
experimentName = `Testing ${timestamp}`;
20-
experimentId = await experimentClient.createExperiment(experimentName);
21-
testIds.push(experimentId);
19+
experimentClient = new ExperimentClient(TRACKING_SERVER_URI);
2220
});
2321

2422
describe('createExperiment', () => {
2523
test('should create an experiment and return the experiment ID', async () => {
26-
const timestamp = Date.now();
27-
const testExperimentId = await experimentClient.createExperiment(
28-
`Test experiment ${timestamp}`
29-
);
24+
const testExperimentId: string = await createTestExperiment();
3025
testIds.push(testExperimentId);
3126
expect(typeof testExperimentId).toBe('string');
3227
expect(testExperimentId).toBeTruthy();
@@ -44,6 +39,8 @@ describe('ExperimentClient', () => {
4439
});
4540

4641
test('should throw error if name is already in use', async () => {
42+
const experimentName: string = `Test experiment ${Date.now()}`;
43+
testIds.push(await experimentClient.createExperiment(experimentName));
4744
await expect(
4845
experimentClient.createExperiment(experimentName)
4946
).rejects.toThrow(ApiError);
@@ -53,18 +50,15 @@ describe('ExperimentClient', () => {
5350
describe('searchExperiment', () => {
5451
beforeAll(async () => {
5552
for (let i = 0; i < 5; i++) {
56-
const num = Math.random().toString().slice(2, 11);
57-
const name = `Search test ${num}`;
58-
const search = await experimentClient.createExperiment(name);
53+
const search: string = await experimentClient.createExperiment(
54+
`Search test ${Date.now()}`
55+
);
5956
testIds.push(search);
6057
}
6158
});
6259

6360
test('should return valid search results', async () => {
64-
const results: {
65-
experiments?: Experiment[];
66-
next_page_token?: string;
67-
} = await experimentClient.searchExperiment(
61+
const results: ExpSearchResults = await experimentClient.searchExperiment(
6862
"name LIKE 'Search test%'",
6963
4
7064
);
@@ -73,26 +67,22 @@ describe('ExperimentClient', () => {
7367
expect(results.next_page_token).toBeDefined();
7468
expect(results.experiments).toHaveLength(4);
7569
results.experiments?.forEach((result) => {
76-
expect(result).toHaveProperty('experiment_id');
77-
expect(result).toHaveProperty('name');
78-
expect(result).toHaveProperty('artifact_location');
79-
expect(result).toHaveProperty('lifecycle_stage');
80-
expect(result).toHaveProperty('last_update_time');
81-
expect(result).toHaveProperty('creation_time');
70+
for (const property of experimentProperties) {
71+
expect(result).toHaveProperty(property);
72+
}
8273
});
8374
expect(typeof results.next_page_token).toBe('string');
8475
});
8576
});
8677

8778
describe('getExperiment', () => {
8879
test('should return experiment information', async () => {
89-
const experiment = await experimentClient.getExperiment(experimentId);
90-
expect(experiment).toHaveProperty('experiment_id');
91-
expect(experiment).toHaveProperty('name');
92-
expect(experiment).toHaveProperty('artifact_location');
93-
expect(experiment).toHaveProperty('lifecycle_stage');
94-
expect(experiment).toHaveProperty('last_update_time');
95-
expect(experiment).toHaveProperty('creation_time');
80+
const expId: string = await createTestExperiment();
81+
const experiment = await experimentClient.getExperiment(expId);
82+
testIds.push(expId);
83+
for (const property of experimentProperties) {
84+
expect(experiment).toHaveProperty(property);
85+
}
9686
});
9787

9888
test('should throw error if experiment ID is missing', async () => {
@@ -103,54 +93,49 @@ describe('ExperimentClient', () => {
10393

10494
describe('getExperimentByName', () => {
10595
test('should return experiment information', async () => {
106-
const experiment = await experimentClient.getExperimentByName(experimentName);
107-
expect(experiment).toHaveProperty('experiment_id');
108-
expect(experiment).toHaveProperty('name');
109-
expect(experiment).toHaveProperty('artifact_location');
110-
expect(experiment).toHaveProperty('lifecycle_stage');
111-
expect(experiment).toHaveProperty('last_update_time');
112-
expect(experiment).toHaveProperty('creation_time');
96+
const name: string = `Test experiment ${Date.now()}`;
97+
testIds.push(await experimentClient.createExperiment(name));
98+
const experiment = await experimentClient.getExperimentByName(name);
99+
for (const property of experimentProperties) {
100+
expect(experiment).toHaveProperty(property);
101+
}
113102
});
114103

115104
test('should throw error if experiment name is missing', async () => {
116105
// @ts-expect-error: testing for missing arguments
117-
await expect(experimentClient.getExperimentByName()).rejects.toThrow(ApiError);
106+
await expect(experimentClient.getExperimentByName()).rejects.toThrow(
107+
ApiError
108+
);
118109
});
119110
});
120111

121112
describe('deleteExperiment', () => {
122113
test('should delete an experiment', async () => {
123-
const num = Math.random().toString().slice(2, 11);
124-
const name = `Test experiment ${num}`;
125-
const idToDelete = await experimentClient.createExperiment(name);
114+
const name: string = `Test experiment ${Date.now()}`;
115+
const idToDelete: string = await experimentClient.createExperiment(name);
126116
await experimentClient.deleteExperiment(idToDelete);
127-
const results: {
128-
experiments?: Experiment[];
129-
next_page_token?: string;
130-
} = await experimentClient.searchExperiment(
131-
`name LIKE '${idToDelete}'`,
117+
const results: ExpSearchResults = await experimentClient.searchExperiment(
118+
`name LIKE '${name}'`,
132119
4
133120
);
134121
expect(results).toEqual({});
135122
});
136123

137124
test('should throw error if invalid experiment ID is passed in', async () => {
138-
await expect(experimentClient.deleteExperiment('invalidExperimentId')).rejects.toThrow(ApiError);
125+
await expect(
126+
experimentClient.deleteExperiment('invalidExperimentId')
127+
).rejects.toThrow(ApiError);
139128
});
140129
});
141130

142131
describe('restoreExperiment', () => {
143132
test('should restore a deleted experiment', async () => {
144-
const num = Math.random().toString().slice(2, 11);
145-
const name = `Test experiment ${num}`;
146-
const idToDelete = await experimentClient.createExperiment(name);
133+
const name: string = `Test experiment ${Date.now()}`;
134+
const idToDelete: string = await experimentClient.createExperiment(name);
147135
testIds.push(idToDelete);
148136
await experimentClient.deleteExperiment(idToDelete);
149137
await experimentClient.restoreExperiment(idToDelete);
150-
const results: {
151-
experiments?: Experiment[];
152-
next_page_token?: string;
153-
} = await experimentClient.searchExperiment(
138+
const results: ExpSearchResults = await experimentClient.searchExperiment(
154139
`name LIKE '${name}'`,
155140
4
156141
);
@@ -159,22 +144,20 @@ describe('ExperimentClient', () => {
159144
});
160145

161146
test('should throw error if invalid experiment ID is passed in', async () => {
162-
await expect(experimentClient.restoreExperiment('invalidExperimentId')).rejects.toThrow(ApiError);
147+
await expect(
148+
experimentClient.restoreExperiment('invalidExperimentId')
149+
).rejects.toThrow(ApiError);
163150
});
164151
});
165152

166153
describe('updateExperiment', () => {
167-
test('should update an experiment\'s name', async () => {
168-
const num = Math.random().toString().slice(2, 11);
169-
const name = `Test experiment ${num}`;
170-
const exp = await experimentClient.createExperiment(name);
154+
test("should update an experiment's name", async () => {
155+
const name: string = `Test experiment ${Date.now()}`;
156+
const exp: string = await experimentClient.createExperiment(name);
171157
testIds.push(exp);
172-
const updatedName = `${name}_UPDATE`
158+
const updatedName: string = `${name}_UPDATE`;
173159
await experimentClient.updateExperiment(exp, updatedName);
174-
const results: {
175-
experiments?: Experiment[];
176-
next_page_token?: string;
177-
} = await experimentClient.searchExperiment(
160+
const results: ExpSearchResults = await experimentClient.searchExperiment(
178161
`name LIKE '${updatedName}'`,
179162
4
180163
);
@@ -184,22 +167,24 @@ describe('ExperimentClient', () => {
184167
});
185168

186169
test('should throw error if invalid experiment ID is passed in', async () => {
187-
await expect(experimentClient.updateExperiment('invalidExperimentId', 'invalidExperimentIdUpdate')).rejects.toThrow(ApiError);
170+
await expect(
171+
experimentClient.updateExperiment(
172+
'invalidExperimentId',
173+
'invalidExperimentIdUpdate'
174+
)
175+
).rejects.toThrow(ApiError);
188176
});
189177
});
190178

191179
describe('setExperimentTag', () => {
192180
test('should set a tag on an experiment', async () => {
193-
const num = Math.random().toString().slice(2, 11);
194-
const name = `Test experiment ${num}`;
195-
const exp = await experimentClient.createExperiment(name);
181+
const date: number = Date.now();
182+
const name: string = `Test experiment ${date}`;
183+
const exp: string = await experimentClient.createExperiment(name);
196184
testIds.push(exp);
197-
await experimentClient.setExperimentTag(exp, 'tag1', `value${num}`);
198-
const results: {
199-
experiments?: Experiment[];
200-
next_page_token?: string;
201-
} = await experimentClient.searchExperiment(
202-
`tags.tag1 = "value${num}"`,
185+
await experimentClient.setExperimentTag(exp, 'tag1', `value${date}`);
186+
const results: ExpSearchResults = await experimentClient.searchExperiment(
187+
`tags.tag1 = "value${date}"`,
203188
4
204189
);
205190
expect(results.experiments).toBeDefined();
@@ -208,16 +193,15 @@ describe('ExperimentClient', () => {
208193
});
209194

210195
test('should throw error if invalid experiment ID is passed in', async () => {
211-
await expect(experimentClient.setExperimentTag('invalidExperimentId', 'tag1', 'value1')).rejects.toThrow(ApiError);
196+
await expect(
197+
experimentClient.setExperimentTag(
198+
'invalidExperimentId',
199+
'tag1',
200+
'value1'
201+
)
202+
).rejects.toThrow(ApiError);
212203
});
213204
});
214205

215-
afterAll(async () => {
216-
while (testIds.length > 0) {
217-
const id = testIds.pop();
218-
if (id) {
219-
await experimentClient.deleteExperiment(id);
220-
}
221-
}
222-
});
206+
afterAll(async () => await deleteTestExperiments(testIds));
223207
});

0 commit comments

Comments
 (0)