Skip to content

Commit 3a54719

Browse files
kartikgupta-dbkartikgupta-db
authored
Add cancelation token to SDK and update tests (#69)
* Add cancelation token and update tests * cleanup * comments * address feedback * cleanup
1 parent 98ad2f6 commit 3a54719

12 files changed

Lines changed: 435 additions & 84 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
"@typescript-eslint/parser": "^5.33.0",
3434
"eslint": "^8.21.0",
3535
"eslint-config-prettier": "^8.5.0",
36+
"ts-mockito": "^2.6.1",
3637
"typescript": "^4.7.4"
3738
}
3839
}

‎packages/databricks-sdk-js/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
"tmp-promise": "^3.0.3",
5353
"ts-loader": "^9.3.1",
5454
"ts-mocha": "^10.0.0",
55+
"ts-mockito": "^2.6.1",
5556
"ts-node": "^10.9.1",
5657
"typescript": "^4.7.4",
5758
"uuid": "^8.3.2"

‎packages/databricks-sdk-js/src/retries/retries.ts‎

Lines changed: 20 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,23 @@ export interface RetriableResult<T> {
99
error?: unknown;
1010
}
1111

12-
const maxJitter = new Time(750, TimeUnits.milliseconds);
13-
const minJitter = new Time(50, TimeUnits.milliseconds);
14-
const maxWaitTime = new Time(10, TimeUnits.seconds);
15-
const defaultTimeout = new Time(10, TimeUnits.minutes);
12+
export class RetryConfigs {
13+
static maxJitter = new Time(750, TimeUnits.milliseconds);
14+
static minJitter = new Time(50, TimeUnits.milliseconds);
15+
static maxWaitTime = new Time(10, TimeUnits.seconds);
16+
static defaultTimeout = new Time(10, TimeUnits.minutes);
1617

17-
function waitTime(attempt: number) {
18-
const jitter = maxJitter
19-
.sub(minJitter)
20-
.multiply(Math.random())
21-
.add(minJitter);
22-
const timeout = new Time(attempt, TimeUnits.seconds).add(jitter);
18+
static waitTime(attempt: number) {
19+
const jitter = RetryConfigs.maxJitter
20+
.sub(RetryConfigs.minJitter)
21+
.multiply(Math.random())
22+
.add(RetryConfigs.minJitter);
23+
const timeout = new Time(attempt, TimeUnits.seconds).add(jitter);
2324

24-
return timeout.gt(maxWaitTime) ? maxWaitTime : timeout;
25+
return timeout.gt(RetryConfigs.maxWaitTime)
26+
? RetryConfigs.maxWaitTime
27+
: timeout;
28+
}
2529
}
2630

2731
interface RetryArgs<T> {
@@ -30,7 +34,7 @@ interface RetryArgs<T> {
3034
}
3135

3236
export default async function retry<T>({
33-
timeout = defaultTimeout,
37+
timeout = RetryConfigs.defaultTimeout,
3438
fn,
3539
}: RetryArgs<T>): Promise<T> {
3640
let attempt = 1;
@@ -60,7 +64,10 @@ export default async function retry<T>({
6064
}
6165

6266
await new Promise((resolve) =>
63-
setTimeout(resolve, waitTime(attempt).toMillSeconds().value)
67+
setTimeout(
68+
resolve,
69+
RetryConfigs.waitTime(attempt).toMillSeconds().value
70+
)
6471
);
6572

6673
attempt += 1;

‎packages/databricks-sdk-js/src/services/Cluster.integ.ts‎

Lines changed: 5 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
/* eslint-disable @typescript-eslint/naming-convention */
22

3-
import {Cluster, ClusterService} from "..";
3+
import {Cluster} from "..";
44
import assert from "assert";
5-
5+
import chai from "chai";
6+
import chaiAsPromised from "chai-as-promised";
67
import {IntegrationTestSetup} from "../test/IntegrationTestSetup";
78

9+
chai.use(chaiAsPromised);
10+
811
describe(__filename, function () {
912
let integSetup: IntegrationTestSetup;
1013

@@ -36,44 +39,4 @@ describe(__filename, function () {
3639
assert(clusterA.id);
3740
assert.equal(clusterA.id, clusterB?.id);
3841
});
39-
40-
// TODO: run tests changing the state of cluster in a seperate job, on a single node
41-
it.skip("calling start on a non terminated state should not throw an error", async () => {
42-
await integSetup.cluster.stop();
43-
await new ClusterService(integSetup.client).start({
44-
cluster_id: integSetup.cluster.id,
45-
});
46-
47-
await integSetup.cluster.refresh();
48-
assert(integSetup.cluster.state !== "RUNNING");
49-
50-
await integSetup.cluster.start();
51-
});
52-
53-
it.skip("should terminate cluster", async () => {
54-
integSetup.cluster.start();
55-
assert.equal(integSetup.cluster.state, "RUNNING");
56-
57-
await integSetup.cluster.stop();
58-
assert.equal(integSetup.cluster.state, "TERMINATED");
59-
60-
integSetup.cluster.start();
61-
});
62-
63-
it.skip("should terminate non running clusters", async () => {
64-
await integSetup.cluster.stop();
65-
assert.equal(integSetup.cluster.state, "TERMINATED");
66-
67-
await new ClusterService(integSetup.client).start({
68-
cluster_id: integSetup.cluster.id,
69-
});
70-
await integSetup.cluster.refresh();
71-
assert.notEqual(integSetup.cluster.state, "RUNNING");
72-
73-
await integSetup.cluster.stop();
74-
assert.equal(integSetup.cluster.state, "TERMINATED");
75-
76-
await integSetup.cluster.start();
77-
assert.equal(integSetup.cluster.state, "RUNNING");
78-
});
7942
});

‎packages/databricks-sdk-js/src/services/Cluster.test.ts‎

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
/* eslint-disable @typescript-eslint/naming-convention */
2+
3+
import {ApiClient, Cluster} from "..";
4+
import chai, {assert} from "chai";
5+
import {mock, when, instance, deepEqual, verify, anything} from "ts-mockito";
6+
import chaiAsPromised from "chai-as-promised";
7+
import Time, {TimeUnits} from "../retries/Time";
8+
import getMockTestCluster from "../test/fixtures/ClusterFixtures";
9+
import {GetClusterResponse} from "../apis/cluster";
10+
import TokenFixture from "../test/fixtures/TokenFixture";
11+
import {RetryConfigs} from "../retries/retries";
12+
13+
chai.use(chaiAsPromised);
14+
15+
describe(__filename, function () {
16+
this.timeout(new Time(10, TimeUnits.minutes).toMillSeconds().value);
17+
18+
let mockedClient: ApiClient;
19+
let mockedCluster: Cluster;
20+
let testClusterDetails: GetClusterResponse;
21+
22+
beforeEach(async () => {
23+
({mockedCluster, mockedClient, testClusterDetails} =
24+
await getMockTestCluster());
25+
26+
RetryConfigs.waitTime = (attempt) => {
27+
return new Time(0, TimeUnits.milliseconds);
28+
};
29+
});
30+
31+
it("calling start on a non terminated state should not throw an error", async () => {
32+
when(
33+
mockedClient.request(
34+
"/api/2.0/clusters/get",
35+
"GET",
36+
deepEqual({
37+
cluster_id: testClusterDetails.cluster_id,
38+
})
39+
)
40+
).thenResolve(
41+
{
42+
...testClusterDetails,
43+
state: "PENDING",
44+
},
45+
{
46+
...testClusterDetails,
47+
state: "PENDING",
48+
},
49+
{
50+
...testClusterDetails,
51+
state: "PENDING",
52+
},
53+
{
54+
...testClusterDetails,
55+
state: "PENDING",
56+
},
57+
{
58+
...testClusterDetails,
59+
state: "PENDING",
60+
},
61+
{
62+
...testClusterDetails,
63+
state: "RUNNING",
64+
}
65+
);
66+
67+
await mockedCluster.refresh();
68+
assert.notEqual(mockedCluster.state, "RUNNING");
69+
70+
await mockedCluster.start();
71+
assert.equal(mockedCluster.state, "RUNNING");
72+
73+
verify(
74+
mockedClient.request(
75+
"/api/2.0/clusters/get",
76+
anything(),
77+
anything()
78+
)
79+
).times(6);
80+
81+
verify(
82+
mockedClient.request(
83+
"/api/2.0/clusters/start",
84+
anything(),
85+
anything()
86+
)
87+
).never();
88+
});
89+
90+
it("should terminate cluster", async () => {
91+
when(
92+
mockedClient.request(
93+
"/api/2.0/clusters/get",
94+
"GET",
95+
deepEqual({
96+
cluster_id: testClusterDetails.cluster_id,
97+
})
98+
)
99+
).thenResolve(
100+
{
101+
...testClusterDetails,
102+
state: "RUNNING",
103+
},
104+
{
105+
...testClusterDetails,
106+
state: "TERMINATING",
107+
},
108+
{
109+
...testClusterDetails,
110+
state: "TERMINATED",
111+
}
112+
);
113+
114+
when(
115+
mockedClient.request(
116+
"/api/2.0/clusters/delete",
117+
"POST",
118+
deepEqual({
119+
cluster_id: testClusterDetails.cluster_id,
120+
})
121+
)
122+
).thenResolve({});
123+
124+
assert.equal(mockedCluster.state, "RUNNING");
125+
126+
await mockedCluster.stop();
127+
assert.equal(mockedCluster.state, "TERMINATED");
128+
129+
verify(
130+
mockedClient.request(
131+
"/api/2.0/clusters/get",
132+
anything(),
133+
anything()
134+
)
135+
).times(3);
136+
137+
verify(
138+
mockedClient.request(
139+
"/api/2.0/clusters/delete",
140+
anything(),
141+
anything()
142+
)
143+
).once();
144+
});
145+
146+
it("should terminate non running clusters", async () => {
147+
when(
148+
mockedClient.request(
149+
"/api/2.0/clusters/get",
150+
"GET",
151+
deepEqual({
152+
cluster_id: testClusterDetails.cluster_id,
153+
})
154+
)
155+
).thenResolve(
156+
{
157+
...testClusterDetails,
158+
state: "PENDING",
159+
},
160+
{
161+
...testClusterDetails,
162+
state: "TERMINATING",
163+
},
164+
{
165+
...testClusterDetails,
166+
state: "TERMINATED",
167+
}
168+
);
169+
170+
await mockedCluster.refresh();
171+
assert.notEqual(mockedCluster.state, "RUNNING");
172+
173+
await mockedCluster.stop();
174+
assert.equal(mockedCluster.state, "TERMINATED");
175+
176+
verify(
177+
mockedClient.request(
178+
"/api/2.0/clusters/get",
179+
anything(),
180+
anything()
181+
)
182+
).times(3);
183+
184+
verify(
185+
mockedClient.request(
186+
"/api/2.0/clusters/delete",
187+
anything(),
188+
anything()
189+
)
190+
).once();
191+
});
192+
193+
it("should cancel cluster start", async () => {
194+
const whenMockGetCluster = when(
195+
mockedClient.request(
196+
"/api/2.0/clusters/get",
197+
"GET",
198+
deepEqual({
199+
cluster_id: testClusterDetails.cluster_id,
200+
})
201+
)
202+
);
203+
204+
whenMockGetCluster.thenResolve({
205+
...testClusterDetails,
206+
state: "PENDING",
207+
});
208+
209+
when(
210+
mockedClient.request(
211+
"/api/2.0/clusters/delete",
212+
"POST",
213+
deepEqual({
214+
cluster_id: testClusterDetails.cluster_id,
215+
})
216+
)
217+
).thenCall(() => {
218+
whenMockGetCluster.thenResolve(
219+
{
220+
...testClusterDetails,
221+
state: "TERMINATING",
222+
},
223+
{
224+
...testClusterDetails,
225+
state: "TERMINATED",
226+
}
227+
);
228+
return {};
229+
});
230+
231+
const token = mock(TokenFixture);
232+
when(token.isCancellationRequested).thenReturn(false, false, true);
233+
//mocked cluster is initially in running state, this gets it to pending state
234+
await mockedCluster.refresh();
235+
236+
assert.equal(mockedCluster.state, "PENDING");
237+
await mockedCluster.start(instance(token));
238+
239+
verify(token.isCancellationRequested).thrice();
240+
assert.equal(mockedCluster.state, "TERMINATED");
241+
});
242+
});

0 commit comments

Comments
 (0)