Skip to content

Commit b940519

Browse files
yiqunzhengyiqunzheng
authored
Merge pull request #88 from oslabs-beta/austin-jest
feat: Jest tests for ExperimentClient and ExperimentManager
2 parents 7e5069d + c86a637 commit b940519

File tree

6 files changed

+424
-238
lines changed

6 files changed

+424
-238
lines changed

‎mlflow/src/utils/interface.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,3 +74,12 @@ export interface CopyRun {
7474
newRunId: string;
7575
targetExperimentId: string;
7676
}
77+
78+
export interface Experiment {
79+
experiment_id: string;
80+
name: string;
81+
artifact_location: string;
82+
lifecycle_stage: string;
83+
last_update_time: string;
84+
creation_time: string;
85+
}

‎mlflow/src/workflows/ExperimentManager.ts

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ class ExperimentManager {
6767
// updateRun to finish it
6868
const latestRun = await this.runClient.updateRun(run_id, 'FINISHED');
6969

70-
return latestRun;
70+
return (latestRun as { run_info: object }).run_info;
7171
} catch (error) {
7272
if (error instanceof ApiError) {
7373
console.error(`API Error (${error.statusCode}): ${error.message}`);
@@ -93,8 +93,8 @@ class ExperimentManager {
9393
*/
9494
async runNewExperiment(
9595
experiment_name: string,
96-
run_name: string,
97-
metrics: Array<{
96+
run_name?: string,
97+
metrics?: Array<{
9898
key: string;
9999
value: number;
100100
timestamp: number;
@@ -133,9 +133,9 @@ class ExperimentManager {
133133
}
134134

135135
// updateRun to finish it
136-
const latest_run = await this.runClient.updateRun(run_id, 'FINISHED');
136+
const latestRun = await this.runClient.updateRun(run_id, 'FINISHED');
137137

138-
return latest_run;
138+
return (latestRun as { run_info: object }).run_info;
139139
} catch (error) {
140140
if (error instanceof ApiError) {
141141
console.error(`API Error (${error.statusCode}): ${error.message}`);

‎mlflow/tests/ExperimentClient.test.ts

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
import { describe, test, expect, beforeAll, afterAll } from '@jest/globals';
2+
import ExperimentClient from '../src/tracking/ExperimentClient';
3+
import { ApiError } from '../src/utils/apiError';
4+
import { Experiment } from '../src/utils/interface';
5+
6+
describe('ExperimentClient', () => {
7+
let experimentClient: ExperimentClient;
8+
let experimentId: string;
9+
let experimentName: string;
10+
const testIds: string[] = [];
11+
12+
beforeAll(async () => {
13+
// Add a small delay to ensure MLflow is fully ready
14+
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);
22+
});
23+
24+
describe('createExperiment', () => {
25+
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+
);
30+
testIds.push(testExperimentId);
31+
expect(typeof testExperimentId).toBe('string');
32+
expect(testExperimentId).toBeTruthy();
33+
});
34+
35+
test('should throw error if name is missing', async () => {
36+
// @ts-expect-error: testing for missing arguments
37+
await expect(experimentClient.createExperiment()).rejects.toThrow(
38+
ApiError
39+
);
40+
// @ts-expect-error: testing for missing arguments
41+
await expect(experimentClient.createExperiment()).rejects.toThrow(
42+
/Error creating experiment from tracking server:/
43+
);
44+
});
45+
46+
test('should throw error if name is already in use', async () => {
47+
await expect(
48+
experimentClient.createExperiment(experimentName)
49+
).rejects.toThrow(ApiError);
50+
});
51+
});
52+
53+
describe('searchExperiment', () => {
54+
beforeAll(async () => {
55+
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);
59+
testIds.push(search);
60+
}
61+
});
62+
63+
test('should return valid search results', async () => {
64+
const results: {
65+
experiments?: Experiment[];
66+
next_page_token?: string;
67+
} = await experimentClient.searchExperiment(
68+
"name LIKE 'Search test%'",
69+
4
70+
);
71+
72+
expect(results.experiments).toBeDefined();
73+
expect(results.next_page_token).toBeDefined();
74+
expect(results.experiments).toHaveLength(4);
75+
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');
82+
});
83+
expect(typeof results.next_page_token).toBe('string');
84+
});
85+
});
86+
87+
describe('getExperiment', () => {
88+
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');
96+
});
97+
98+
test('should throw error if experiment ID is missing', async () => {
99+
// @ts-expect-error: testing for missing arguments
100+
await expect(experimentClient.getExperiment()).rejects.toThrow(ApiError);
101+
});
102+
});
103+
104+
describe('getExperimentByName', () => {
105+
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');
113+
});
114+
115+
test('should throw error if experiment name is missing', async () => {
116+
// @ts-expect-error: testing for missing arguments
117+
await expect(experimentClient.getExperimentByName()).rejects.toThrow(ApiError);
118+
});
119+
});
120+
121+
describe('deleteExperiment', () => {
122+
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);
126+
await experimentClient.deleteExperiment(idToDelete);
127+
const results: {
128+
experiments?: Experiment[];
129+
next_page_token?: string;
130+
} = await experimentClient.searchExperiment(
131+
`name LIKE '${idToDelete}'`,
132+
4
133+
);
134+
expect(results).toEqual({});
135+
});
136+
137+
test('should throw error if invalid experiment ID is passed in', async () => {
138+
await expect(experimentClient.deleteExperiment('invalidExperimentId')).rejects.toThrow(ApiError);
139+
});
140+
});
141+
142+
describe('restoreExperiment', () => {
143+
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);
147+
testIds.push(idToDelete);
148+
await experimentClient.deleteExperiment(idToDelete);
149+
await experimentClient.restoreExperiment(idToDelete);
150+
const results: {
151+
experiments?: Experiment[];
152+
next_page_token?: string;
153+
} = await experimentClient.searchExperiment(
154+
`name LIKE '${name}'`,
155+
4
156+
);
157+
expect(results.experiments).toBeDefined();
158+
expect(results.experiments).toHaveLength(1);
159+
});
160+
161+
test('should throw error if invalid experiment ID is passed in', async () => {
162+
await expect(experimentClient.restoreExperiment('invalidExperimentId')).rejects.toThrow(ApiError);
163+
});
164+
});
165+
166+
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);
171+
testIds.push(exp);
172+
const updatedName = `${name}_UPDATE`
173+
await experimentClient.updateExperiment(exp, updatedName);
174+
const results: {
175+
experiments?: Experiment[];
176+
next_page_token?: string;
177+
} = await experimentClient.searchExperiment(
178+
`name LIKE '${updatedName}'`,
179+
4
180+
);
181+
expect(results.experiments).toBeDefined();
182+
expect(results.experiments).toHaveLength(1);
183+
expect(results.experiments?.[0].experiment_id).toBe(exp);
184+
});
185+
186+
test('should throw error if invalid experiment ID is passed in', async () => {
187+
await expect(experimentClient.updateExperiment('invalidExperimentId', 'invalidExperimentIdUpdate')).rejects.toThrow(ApiError);
188+
});
189+
});
190+
191+
describe('setExperimentTag', () => {
192+
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);
196+
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}"`,
203+
4
204+
);
205+
expect(results.experiments).toBeDefined();
206+
expect(results.experiments).toHaveLength(1);
207+
expect(results.experiments?.[0].experiment_id).toBe(exp);
208+
});
209+
210+
test('should throw error if invalid experiment ID is passed in', async () => {
211+
await expect(experimentClient.setExperimentTag('invalidExperimentId', 'tag1', 'value1')).rejects.toThrow(ApiError);
212+
});
213+
});
214+
215+
afterAll(async () => {
216+
while (testIds.length > 0) {
217+
const id = testIds.pop();
218+
if (id) {
219+
await experimentClient.deleteExperiment(id);
220+
}
221+
}
222+
});
223+
});

‎mlflow/tests/ExperimentClientTestFile.ts

Lines changed: 0 additions & 123 deletions
This file was deleted.

0 commit comments

Comments
 (0)