Skip to content

Commit bea0807

Browse files
author
John Schulz
committed
Refactor to make more testable. Add more tests.
1 parent b956c61 commit bea0807

2 files changed

Lines changed: 202 additions & 67 deletions

File tree

‎x-pack/plugins/ingest_manager/server/routes/limited_concurrency.test.ts‎

Lines changed: 171 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,11 @@
55
*/
66

77
import { coreMock, httpServerMock, httpServiceMock } from 'src/core/server/mocks';
8-
import { registerLimitedConcurrencyRoutes } from './limited_concurrency';
8+
import {
9+
createLimitedPreAuthHandler,
10+
isLimitedRoute,
11+
registerLimitedConcurrencyRoutes,
12+
} from './limited_concurrency';
913
import { IngestManagerConfigType } from '../index';
1014

1115
describe('registerLimitedConcurrencyRoutes', () => {
@@ -35,70 +39,184 @@ describe('registerLimitedConcurrencyRoutes', () => {
3539
});
3640

3741
describe('preAuthHandler', () => {
38-
test(`it ignores routes which don't have the correct tag`, async () => {
39-
const routerMock = coreMock.createSetup().http.createRouter();
40-
const handlerMock = jest.fn();
41-
routerMock.get(
42-
{
43-
path: '/api/foo',
44-
validate: false,
45-
// options: { tags: ['ingest:limited-concurrency'] },
46-
},
47-
handlerMock
48-
);
42+
test(`ignores routes when !isMatch`, async () => {
43+
const mockMaxCounter = {
44+
increase: jest.fn(),
45+
decrease: jest.fn(),
46+
lessThanMax: jest.fn(),
47+
};
48+
const preAuthHandler = createLimitedPreAuthHandler({
49+
isMatch: jest.fn().mockImplementation(() => false),
50+
maxCounter: mockMaxCounter,
51+
});
4952

50-
const mockSetup = coreMock.createSetup();
51-
const mockConfig = { fleet: { maxConcurrentConnections: 1 } } as IngestManagerConfigType;
52-
registerLimitedConcurrencyRoutes(mockSetup, mockConfig);
53+
const mockRequest = httpServerMock.createKibanaRequest({
54+
path: '/no/match',
55+
});
56+
const mockResponse = httpServerMock.createResponseFactory();
57+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
58+
59+
// @ts-ignore error re: mockPreAuthToolkit return type
60+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
61+
62+
expect(mockMaxCounter.increase).not.toHaveBeenCalled();
63+
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
64+
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
65+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
66+
});
67+
68+
test(`ignores routes which don't have the correct tag`, async () => {
69+
const mockMaxCounter = {
70+
increase: jest.fn(),
71+
decrease: jest.fn(),
72+
lessThanMax: jest.fn(),
73+
};
74+
const preAuthHandler = createLimitedPreAuthHandler({
75+
isMatch: isLimitedRoute,
76+
maxCounter: mockMaxCounter,
77+
});
78+
79+
const mockRequest = httpServerMock.createKibanaRequest({
80+
path: '/no/match',
81+
});
82+
const mockResponse = httpServerMock.createResponseFactory();
83+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
5384

54-
const [[preAuthHandler]] = mockSetup.http.registerOnPreAuth.mock.calls;
85+
// @ts-ignore error re: mockPreAuthToolkit return type
86+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
87+
88+
expect(mockMaxCounter.increase).not.toHaveBeenCalled();
89+
expect(mockMaxCounter.decrease).not.toHaveBeenCalled();
90+
expect(mockMaxCounter.lessThanMax).not.toHaveBeenCalled();
91+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
92+
});
93+
94+
test(`processes routes which have the correct tag`, async () => {
95+
const mockMaxCounter = {
96+
increase: jest.fn(),
97+
decrease: jest.fn(),
98+
lessThanMax: jest.fn().mockImplementation(() => true),
99+
};
100+
const preAuthHandler = createLimitedPreAuthHandler({
101+
isMatch: isLimitedRoute,
102+
maxCounter: mockMaxCounter,
103+
});
55104

56105
const mockRequest = httpServerMock.createKibanaRequest({
57-
method: 'get',
58-
path: '/api/foo',
59-
// routeTags: ['ingest:limited-concurrency'],
106+
path: '/should/match',
107+
routeTags: ['ingest:limited-concurrency'],
60108
});
61109
const mockResponse = httpServerMock.createResponseFactory();
62110
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
63111

64112
// @ts-ignore error re: mockPreAuthToolkit return type
65113
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
66-
const [routeConfig, routeHandler] = routerMock.get.mock.calls[0];
67-
// @ts-ignore error re: missing iterator
68-
// const [routeConfig, routeHandler] = routerMock.get.mock.calls.find(
69-
// ([{ path }]) => path === '/api/foo'
70-
// );
71-
console.log({ routeConfig, routeHandler });
72-
expect(mockResponse.notFound).not.toHaveBeenCalled();
114+
115+
// will call lessThanMax because isMatch succeeds
116+
expect(mockMaxCounter.lessThanMax).toHaveBeenCalledTimes(1);
117+
// will not error because lessThanMax is true
118+
expect(mockResponse.customError).not.toHaveBeenCalled();
73119
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
120+
});
121+
122+
test(`updates the counter when isMatch & lessThanMax`, async () => {
123+
const mockMaxCounter = {
124+
increase: jest.fn(),
125+
decrease: jest.fn(),
126+
lessThanMax: jest.fn().mockImplementation(() => true),
127+
};
128+
const preAuthHandler = createLimitedPreAuthHandler({
129+
isMatch: jest.fn().mockImplementation(() => true),
130+
maxCounter: mockMaxCounter,
131+
});
74132

75-
expect(handlerMock).toHaveBeenCalled();
133+
const mockRequest = httpServerMock.createKibanaRequest();
134+
const mockResponse = httpServerMock.createResponseFactory();
135+
const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
136+
137+
// @ts-ignore error re: mockPreAuthToolkit return type
138+
await preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit);
139+
140+
expect(mockMaxCounter.increase).toHaveBeenCalled();
141+
// expect(mockMaxCounter.decrease).toHaveBeenCalled();
142+
expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
76143
});
77144

78-
// test(`it processes routes which have the correct tag`, async () => {
79-
// const mockSetup = coreMock.createSetup();
80-
// const mockConfig = { fleet: { maxConcurrentConnections: 1 } } as IngestManagerConfigType;
81-
// registerLimitedConcurrencyRoutes(mockSetup, mockConfig);
82-
83-
// const [[preAuthHandler]] = mockSetup.http.registerOnPreAuth.mock.calls;
84-
85-
// const mockRequest = httpServerMock.createKibanaRequest({
86-
// method: 'get',
87-
// path: '/api/foo',
88-
// routeTags: ['ingest:limited-concurrency'],
89-
// });
90-
// const mockResponse = httpServerMock.createResponseFactory();
91-
// const mockPreAuthToolkit = httpServiceMock.createOnPreAuthToolkit();
92-
93-
// const responses = await Promise.all([
94-
// preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit),
95-
// preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit),
96-
// preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit),
97-
// preAuthHandler(mockRequest, mockResponse, mockPreAuthToolkit),
98-
// ]);
99-
// console.log('first res', responses[0]);
100-
// console.log('last res', responses[3]);
101-
// expect(mockResponse.notFound).not.toHaveBeenCalled();
102-
// expect(mockPreAuthToolkit.next).toHaveBeenCalledTimes(1);
103-
// });
145+
test(`lessThanMax ? next : error`, async () => {
146+
const mockMaxCounter = {
147+
increase: jest.fn(),
148+
decrease: jest.fn(),
149+
lessThanMax: jest
150+
.fn()
151+
// call 1
152+
.mockImplementationOnce(() => true)
153+
// calls 2, 3, 4
154+
.mockImplementationOnce(() => false)
155+
.mockImplementationOnce(() => false)
156+
.mockImplementationOnce(() => false)
157+
// calls 5+
158+
.mockImplementationOnce(() => true)
159+
.mockImplementation(() => true),
160+
};
161+
162+
const preAuthHandler = createLimitedPreAuthHandler({
163+
isMatch: isLimitedRoute,
164+
maxCounter: mockMaxCounter,
165+
});
166+
167+
function makeRequestExpectNext() {
168+
const request = httpServerMock.createKibanaRequest({
169+
path: '/should/match/',
170+
routeTags: ['ingest:limited-concurrency'],
171+
});
172+
const response = httpServerMock.createResponseFactory();
173+
const toolkit = httpServiceMock.createOnPreAuthToolkit();
174+
175+
// @ts-ignore error re: mockPreAuthToolkit return type
176+
preAuthHandler(request, response, toolkit);
177+
expect(toolkit.next).toHaveBeenCalledTimes(1);
178+
expect(response.customError).not.toHaveBeenCalled();
179+
}
180+
181+
function makeRequestExpectError() {
182+
const request = httpServerMock.createKibanaRequest({
183+
path: '/should/match/',
184+
routeTags: ['ingest:limited-concurrency'],
185+
});
186+
const response = httpServerMock.createResponseFactory();
187+
const toolkit = httpServiceMock.createOnPreAuthToolkit();
188+
189+
// @ts-ignore error re: mockPreAuthToolkit return type
190+
preAuthHandler(request, response, toolkit);
191+
expect(toolkit.next).not.toHaveBeenCalled();
192+
expect(response.customError).toHaveBeenCalledTimes(1);
193+
expect(response.customError).toHaveBeenCalledWith({
194+
statusCode: 429,
195+
body: 'Too Many Requests',
196+
});
197+
}
198+
199+
// request 1 succeeds
200+
makeRequestExpectNext();
201+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(1);
202+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(1);
203+
204+
// requests 2, 3, 4 fail
205+
makeRequestExpectError();
206+
makeRequestExpectError();
207+
makeRequestExpectError();
208+
209+
// requests 5+ succeed
210+
makeRequestExpectNext();
211+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(2);
212+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(2);
213+
214+
makeRequestExpectNext();
215+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(3);
216+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(3);
217+
218+
makeRequestExpectNext();
219+
expect(mockMaxCounter.increase).toHaveBeenCalledTimes(4);
220+
// expect(mockMaxCounter.decrease).toHaveBeenCalledTimes(4);
221+
});
104222
});

‎x-pack/plugins/ingest_manager/server/routes/limited_concurrency.ts‎

Lines changed: 31 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ import {
1212
} from 'kibana/server';
1313
import { LIMITED_CONCURRENCY_ROUTE_TAG } from '../../common';
1414
import { IngestManagerConfigType } from '../index';
15-
class MaxCounter {
15+
16+
export class MaxCounter {
1617
constructor(private readonly max: number = 1) {}
1718
private counter = 0;
1819
valueOf() {
@@ -33,40 +34,56 @@ class MaxCounter {
3334
}
3435
}
3536

36-
function shouldHandleRequest(request: KibanaRequest) {
37+
export type IMaxCounter = Pick<MaxCounter, 'increase' | 'decrease' | 'lessThanMax'>;
38+
39+
export function isLimitedRoute(request: KibanaRequest) {
3740
const tags = request.route.options.tags;
38-
return tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
41+
return !!tags.includes(LIMITED_CONCURRENCY_ROUTE_TAG);
3942
}
4043

41-
export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
42-
const max = config.fleet.maxConcurrentConnections;
43-
if (!max) return;
44-
45-
const counter = new MaxCounter(max);
46-
core.http.registerOnPreAuth(function preAuthHandler(
44+
export function createLimitedPreAuthHandler({
45+
isMatch,
46+
maxCounter,
47+
}: {
48+
isMatch: (request: KibanaRequest) => boolean;
49+
maxCounter: IMaxCounter;
50+
}) {
51+
return function preAuthHandler(
4752
request: KibanaRequest,
4853
response: LifecycleResponseFactory,
4954
toolkit: OnPreAuthToolkit
5055
) {
51-
if (!shouldHandleRequest(request)) {
56+
if (!isMatch(request)) {
5257
return toolkit.next();
5358
}
5459

55-
if (!counter.lessThanMax()) {
60+
if (!maxCounter.lessThanMax()) {
5661
return response.customError({
5762
body: 'Too Many Requests',
5863
statusCode: 429,
5964
});
6065
}
6166

62-
counter.increase();
67+
maxCounter.increase();
6368

6469
// requests.events.aborted$ has a bug (but has test which explicitly verifies) where it's fired even when the request completes
6570
// https://github.com/elastic/kibana/pull/70495#issuecomment-656288766
6671
request.events.aborted$.toPromise().then(() => {
67-
counter.decrease();
72+
maxCounter.decrease();
6873
});
6974

7075
return toolkit.next();
71-
});
76+
};
77+
}
78+
79+
export function registerLimitedConcurrencyRoutes(core: CoreSetup, config: IngestManagerConfigType) {
80+
const max = config.fleet.maxConcurrentConnections;
81+
if (!max) return;
82+
83+
core.http.registerOnPreAuth(
84+
createLimitedPreAuthHandler({
85+
isMatch: isLimitedRoute,
86+
maxCounter: new MaxCounter(max),
87+
})
88+
);
7289
}

0 commit comments

Comments
 (0)