Skip to content

Commit

Permalink
feat(instance_launcher): basic tests for launcher logic (#172)
Browse files Browse the repository at this point in the history
* feat(instance_launcher): basic tests for launcher logic

* further launcher logic tests

* log protected status during scaleUp action

* fix: scale protection for new instances check backwards

* better tests for redis

* further tests

* fix tests, update redis logging

* ensure await before negating value
  • Loading branch information
aaronkvanmeerten authored Jan 6, 2025
1 parent b8e10a0 commit 06220b2
Show file tree
Hide file tree
Showing 12 changed files with 738 additions and 32 deletions.
12 changes: 7 additions & 5 deletions src/cloud_manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ export interface CloudManagerOptions extends CloudInstanceManagerSelectorOptions
shutdownManager: ShutdownManager;
instanceTracker: InstanceTracker;
audit: Audit;
cloudInstanceManagerSelector?: CloudInstanceManagerSelector;
}

export interface CloudInstance {
Expand All @@ -34,14 +35,15 @@ export default class CloudManager {
constructor(options: CloudManagerOptions) {
this.isDryRun = options.isDryRun;

this.cloudInstanceManagerSelector = new CloudInstanceManagerSelector(options);
if (options.cloudInstanceManagerSelector) {
this.cloudInstanceManagerSelector = options.cloudInstanceManagerSelector;
} else {
this.cloudInstanceManagerSelector = new CloudInstanceManagerSelector(options);
}

this.instanceTracker = options.instanceTracker;
this.shutdownManager = options.shutdownManager;
this.audit = options.audit;

this.scaleUp = this.scaleUp.bind(this);
this.scaleDown = this.scaleDown.bind(this);
}

async recordLaunch(
Expand Down Expand Up @@ -88,7 +90,7 @@ export default class CloudManager {
isScaleDownProtected: boolean,
): Promise<number> {
const groupName = group.name;
ctx.logger.info('[CloudManager] Scaling up', { groupName, quantity });
ctx.logger.info('[CloudManager] Scaling up', { scaleUp: { groupName, quantity, isScaleDownProtected } });

const instanceManager = this.cloudInstanceManagerSelector.selectInstanceManager(group.cloud);
if (!instanceManager) {
Expand Down
2 changes: 1 addition & 1 deletion src/handlers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -313,7 +313,7 @@ class Handlers {
// found the group, so find the instances and act upon them
// build the list of current instances
const currentInventory = await this.instanceTracker.trimCurrent(req.context, req.params.name);
const instances = this.instanceTracker.mapToInstanceDetails(currentInventory);
const instances = InstanceTracker.mapToInstanceDetails(currentInventory);
// set their reconfigure status to the current date
try {
await this.reconfigureManager.setReconfigureDate(req.context, instances);
Expand Down
10 changes: 7 additions & 3 deletions src/instance_group.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,8 +129,9 @@ export default class InstanceGroupManager {
await this.instanceStore.deleteInstanceGroup(ctx, groupName);
}

// only allow autoscaling if the autoscale grace period has expired
async allowAutoscaling(ctx: Context, group: string): Promise<boolean> {
return this.instanceStore.checkValue(ctx, `autoScaleGracePeriod:${group}`);
return !(await this.instanceStore.checkValue(ctx, `autoScaleGracePeriod:${group}`));
}

async setAutoScaleGracePeriod(ctx: Context, group: InstanceGroup): Promise<boolean> {
Expand All @@ -144,20 +145,23 @@ export default class InstanceGroupManager {
return this.setValue(ctx, `isScaleDownProtected:${group.name}`, group.protectedTTLSec);
}

// only show scale protection if value is set
async isScaleDownProtected(ctx: Context, group: string): Promise<boolean> {
return this.instanceStore.checkValue(ctx, `isScaleDownProtected:${group}`);
}

// only allow group jobs if the grace period has expired
async isGroupJobsCreationAllowed(ctx: Context): Promise<boolean> {
return this.instanceStore.checkValue(ctx, 'groupJobsCreationGracePeriod');
return !(await this.instanceStore.checkValue(ctx, 'groupJobsCreationGracePeriod'));
}

async setGroupJobsCreationGracePeriod(ctx: Context): Promise<boolean> {
return this.setValue(ctx, `groupJobsCreationGracePeriod`, this.processingIntervalSeconds);
}

// only allow sanity jobs if the grace period has expired
async isSanityJobsCreationAllowed(ctx: Context): Promise<boolean> {
return this.instanceStore.checkValue(ctx, 'sanityJobsCreationGracePeriod');
return !(await this.instanceStore.checkValue(ctx, 'sanityJobsCreationGracePeriod'));
}

async setSanityJobsCreationGracePeriod(ctx: Context): Promise<boolean> {
Expand Down
12 changes: 5 additions & 7 deletions src/instance_launcher.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,8 +56,6 @@ export default class InstanceLauncher {
if (options.maxThrottleThreshold) {
this.maxThrottleThreshold = options.maxThrottleThreshold;
}

this.launchOrShutdownInstancesByGroup = this.launchOrShutdownInstancesByGroup.bind(this);
}

async launchOrShutdownInstancesByGroup(ctx: Context, groupName: string): Promise<boolean> {
Expand Down Expand Up @@ -350,7 +348,7 @@ export default class InstanceLauncher {
instanceState.status.provisioning == true
);
});
return this.instanceTracker.mapToInstanceDetails(states);
return InstanceTracker.mapToInstanceDetails(states);
}

private getRunningInstances(instanceStates: InstanceState[]): InstanceDetails[] {
Expand All @@ -364,7 +362,7 @@ export default class InstanceLauncher {
instanceState.status.provisioning == false
);
});
return this.instanceTracker.mapToInstanceDetails(states);
return InstanceTracker.mapToInstanceDetails(states);
}

private getAvailableJibris(instanceStates: InstanceState[]): InstanceDetails[] {
Expand All @@ -373,7 +371,7 @@ export default class InstanceLauncher {
instanceState.status.jibriStatus && instanceState.status.jibriStatus.busyStatus == JibriStatusState.Idle
);
});
return this.instanceTracker.mapToInstanceDetails(states);
return InstanceTracker.mapToInstanceDetails(states);
}

private getExpiredJibris(instanceStates: InstanceState[]): InstanceDetails[] {
Expand All @@ -383,7 +381,7 @@ export default class InstanceLauncher {
instanceState.status.jibriStatus.busyStatus == JibriStatusState.Expired
);
});
return this.instanceTracker.mapToInstanceDetails(states);
return InstanceTracker.mapToInstanceDetails(states);
}

private getBusyJibris(instanceStates: InstanceState[]): InstanceDetails[] {
Expand All @@ -392,6 +390,6 @@ export default class InstanceLauncher {
instanceState.status.jibriStatus && instanceState.status.jibriStatus.busyStatus == JibriStatusState.Busy
);
});
return this.instanceTracker.mapToInstanceDetails(states);
return InstanceTracker.mapToInstanceDetails(states);
}
}
2 changes: 1 addition & 1 deletion src/instance_tracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -464,7 +464,7 @@ export class InstanceTracker {
return states.filter((_, index) => !statesShutdownStatus[index] && !shutdownConfirmations[index]);
}

mapToInstanceDetails(states: InstanceState[]): InstanceDetails[] {
static mapToInstanceDetails(states: InstanceState[]): InstanceDetails[] {
return states.map((response) => {
return <InstanceDetails>{
instanceId: response.instanceId,
Expand Down
36 changes: 21 additions & 15 deletions src/redis.ts
Original file line number Diff line number Diff line change
Expand Up @@ -400,7 +400,7 @@ export default class RedisStore implements MetricsStore, InstanceStore {
return res;
}

async existsAtLeastOneGroup(_ctx: Context): Promise<boolean> {
async existsAtLeastOneGroup(ctx: Context): Promise<boolean> {
let cursor = '0';
do {
const result = await this.redisClient.hscan(
Expand All @@ -411,29 +411,34 @@ export default class RedisStore implements MetricsStore, InstanceStore {
'COUNT',
this.redisScanCount,
);
cursor = result[0];
if (result[1].length > 0) {
const pipeline = this.redisClient.pipeline();
result[1].forEach((key: string) => {
pipeline.hget(this.GROUPS_HASH_NAME, key);
});
if (result) {
cursor = result[0];
if (result[1].length > 0) {
const pipeline = this.redisClient.pipeline();
result[1].forEach((key: string) => {
pipeline.hget(this.GROUPS_HASH_NAME, key);
});

const items = await pipeline.exec();
if (items) {
if (items.length > 0) {
return true;
const items = await pipeline.exec();
if (items) {
if (items.length > 0) {
return true;
}
} else {
return false;
}
} else {
return false;
}
} else {
ctx.logger.error('Error scanning groups for existsAtLeastOneGroup');
return false;
}
} while (cursor != '0');

return false;
}

async upsertInstanceGroup(ctx: Context, group: InstanceGroup): Promise<boolean> {
ctx.logger.info(`Storing ${group.name}`);
ctx.logger.info(`Storing ${group.name}`, { group });
await this.redisClient.hset(this.GROUPS_HASH_NAME, group.name, JSON.stringify(group));
return true;
}
Expand Down Expand Up @@ -507,8 +512,9 @@ export default class RedisStore implements MetricsStore, InstanceStore {

async checkValue(_ctx: Context, key: string): Promise<boolean> {
const result = await this.redisClient.get(key);
return !(result !== null && result.length > 0);
return result !== null && result.length > 0;
}

async setValue(_ctx: Context, key: string, value: string, ttl: number): Promise<boolean> {
const result = await this.redisClient.set(key, value, 'EX', ttl);
if (result !== 'OK') {
Expand Down
119 changes: 119 additions & 0 deletions src/test/cloud_manager.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,119 @@
/* eslint-disable @typescript-eslint/ban-ts-comment */
// @ts-nocheck

import assert from 'node:assert';
import test, { afterEach, describe, mock } from 'node:test';

import CloudManager from '../cloud_manager';

function log(msg, obj) {
console.log(msg, JSON.stringify(obj));
}

function initContext(): Context {
return {
logger: {
info: mock.fn(log),
debug: mock.fn(log),
error: mock.fn(log),
warn: mock.fn(log),
},
};
}

describe('CloudManager', () => {
let context = initContext();

const shutdownManager = {
setScaleDownProtected: mock.fn(() => true),
areScaleDownProtected: mock.fn((arr) => arr.map(() => false)),
};

const instanceTracker = {
track: mock.fn(),
trimCurrent: mock.fn(),
};

const audit = {
saveLaunchEvent: mock.fn(),
};

const cloudInstanceManager = {
getInstances: mock.fn(() => []),
launchInstances: mock.fn((_ctx, _group, _cur, count) => {
for (let i = 0; i < count; i++) {
return [`mock-instance-${i}`];
}
}),
};

const cloudInstanceManagerSelector = {
selectInstanceManager: mock.fn(() => cloudInstanceManager),
};

const cloudManager = new CloudManager({
shutdownManager,
instanceTracker,
audit,
cloudInstanceManagerSelector,
cloudProviders: ['mock'],
isDryRun: false,
});

afterEach(() => {
context = initContext();
cloudInstanceManager.launchInstances.mock.resetCalls();
instanceTracker.track.mock.resetCalls();
});

describe('cloudManager', () => {
test('should return empty for a group with no instances', async () => {
const result = await cloudManager.getInstances(context, 'group');
assert.deepEqual(result, [], 'no instances');
});

test('should return an item for a group with mock instances', async () => {
const instance = { instanceId: 'mock-instance', cloudStatus: 'running', displayName: 'mockery' };
cloudInstanceManager.getInstances.mock.mockImplementation(() => [instance]);
const result = await cloudManager.getInstances(context, 'group');
assert.deepEqual(result, [instance], 'mock instance');
});

test('scaleUp should return success for mock group', async () => {
const group = { name: 'mock-group', cloud: 'mock', type: 'mock' };
const currentCount = 0;
const quantity = 1;
const isScaleDownProtected = false;
const result = await cloudManager.scaleUp(context, group, currentCount, quantity, isScaleDownProtected);
assert.equal(result, 1, 'scale up success');
assert.equal(cloudInstanceManager.launchInstances.mock.calls.length, 1, 'launch instances called');
assert.equal(instanceTracker.track.mock.calls.length, 1, 'newly launched instance tracked');
assert.equal(
shutdownManager.setScaleDownProtected.mock.calls.length,
0,
'launched instance is not scale protected',
);
});

test('scaleUp should return success for mock group, including protected instance call', async () => {
const group = { name: 'mock-group-protected', cloud: 'mock', type: 'mock' };
const currentCount = 0;
const quantity = 1;
const isScaleDownProtected = true;
const result = await cloudManager.scaleUp(context, group, currentCount, quantity, isScaleDownProtected);
assert.equal(result, 1, 'scale up success');
assert.equal(cloudInstanceManager.launchInstances.mock.calls.length, 1, 'launch instances called');
assert.equal(instanceTracker.track.mock.calls.length, 1, 'newly launched instance tracked');
assert.equal(
shutdownManager.setScaleDownProtected.mock.calls.length,
1,
'launched instance is scale protected',
);
assert.equal(
shutdownManager.setScaleDownProtected.mock.calls[0].arguments[2],
instanceTracker.track.mock.calls[0].arguments[1].instanceId,
'protected instance id should match tracked instance id',
);
});
});
});
Loading

0 comments on commit 06220b2

Please sign in to comment.