Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

821 Spike Testing Script #845

Merged
merged 9 commits into from
Feb 29, 2024
40 changes: 2 additions & 38 deletions k6/load.js
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import { check, sleep } from 'k6';
import exec from 'k6/execution';
import http from 'k6/http';
import { postWithSessionUpdate } from './requestFunctions.js';

export const options = {
//vus and duration can be changed depending on what simulation needed running
Expand All @@ -13,40 +11,6 @@ export const options = {
},
};

const baseUrl = 'http://localhost:3001';
const cookieName = 'prompt-injection.sid';
const vuCookieJar = (() => {
const cookieJars = {};
return {
get: (id) => cookieJars[id],
set: (id, jar) => (cookieJars[id] = jar),
};
})();

export default () => {
// Use same jar for every iteration of same VU! k6 doesn't do this for us :(
const vuID = exec.vu.idInTest;
let jar = vuCookieJar.get(vuID);
if (!jar) {
jar = http.cookieJar();
vuCookieJar.set(vuID, jar);
}
let originalCookie = (jar.cookiesForURL(baseUrl)[cookieName] || [])[0];

const data = { infoMessage: 'Hi', chatMessageType: 'LEVEL_INFO', level: 3 };
const response = http.post(`${baseUrl}/test/load`, JSON.stringify(data), {
headers: { 'Content-Type': 'application/json' },
jar,
});
// Expecting cookie to match original, OR first-time be added to the jar
const expectedCookie =
originalCookie || jar.cookiesForURL(baseUrl)[cookieName][0];
check(response, {
'response code was 200': (response) => response.status === 200,
'cookie was preserved': (response) =>
response.cookies[cookieName].length === 1 &&
response.cookies[cookieName][0].value === expectedCookie,
});

sleep(1);
postWithSessionUpdate();
};
41 changes: 41 additions & 0 deletions k6/requestFunctions.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
import { check, sleep } from 'k6';
import exec from 'k6/execution';
import http from 'k6/http';

const baseUrl = 'http://localhost:3001';
const cookieName = 'prompt-injection.sid';
const vuCookieJar = (() => {
const cookieJars = {};
return {
get: (id) => cookieJars[id],
set: (id, jar) => (cookieJars[id] = jar),
};
})();

export function postWithSessionUpdate() {
// Use same jar for every iteration of same VU! k6 doesn't do this for us :(
const vuID = exec.vu.idInTest;
let jar = vuCookieJar.get(vuID);
if (!jar) {
jar = http.cookieJar();
vuCookieJar.set(vuID, jar);
}
let originalCookie = (jar.cookiesForURL(baseUrl)[cookieName] || [])[0];

const data = { infoMessage: 'Hi', chatMessageType: 'LEVEL_INFO', level: 3 };
const response = http.post(`${baseUrl}/test/load`, JSON.stringify(data), {
headers: { 'Content-Type': 'application/json' },
jar,
});
// Expecting cookie to match original, OR first-time be added to the jar
const expectedCookie =
originalCookie || jar.cookiesForURL(baseUrl)[cookieName][0];
check(response, {
'response code was 200': (response) => response.status === 200,
'cookie was preserved': (response) =>
response.cookies[cookieName].length === 1 &&
response.cookies[cookieName][0].value === expectedCookie,
});

sleep(1);
}
13 changes: 13 additions & 0 deletions k6/spike.js
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
import { postWithSessionUpdate } from './requestFunctions.js';

export const options = {
// Key configurations for spike in this section
stages: [
{ duration: '2ms', target: 2000 }, // fast ramp-up to a high point
{ duration: '1m', target: 0 }, // quick ramp-down to 0 users
],
};

export default () => {
kclark-scottlogic marked this conversation as resolved.
Show resolved Hide resolved
postWithSessionUpdate();
};
Loading