wehub-resource-sync
441 lines
17 KiB
TypeScript
441 lines
17 KiB
TypeScript
import { Configuration, EventType, KeyValueStore, Session, SessionPool } from '@crawlee/core';
|
|
import { entries } from '@crawlee/utils';
|
|
|
|
import { Log } from '@apify/log';
|
|
|
|
import { MemoryStorageEmulator } from '../../shared/MemoryStorageEmulator';
|
|
|
|
describe('SessionPool - testing session pool', () => {
|
|
let sessionPool: SessionPool;
|
|
const localStorageEmulator = new MemoryStorageEmulator();
|
|
const events = Configuration.getEventManager();
|
|
|
|
beforeEach(async () => {
|
|
await localStorageEmulator.init();
|
|
sessionPool = await SessionPool.open();
|
|
});
|
|
|
|
afterEach(async () => {
|
|
events.off(EventType.PERSIST_STATE);
|
|
});
|
|
|
|
afterAll(async () => {
|
|
await localStorageEmulator.destroy();
|
|
});
|
|
|
|
test('should initialize with default values for first time', async () => {
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toBeDefined();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.maxPoolSize).toBeDefined();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessionOptions).toBeDefined();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.persistStateKey).toBeDefined();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.createSessionFunction).toEqual(sessionPool._defaultCreateSessionFunction);
|
|
});
|
|
|
|
test('should override default values', async () => {
|
|
const opts = {
|
|
maxPoolSize: 3000,
|
|
sessionOptions: {
|
|
maxAgeSecs: 100,
|
|
maxUsageCount: 1,
|
|
},
|
|
|
|
persistStateKeyValueStoreId: 'TEST',
|
|
persistStateKey: 'SESSION_POOL_STATE2',
|
|
|
|
createSessionFunction: () => ({}) as never,
|
|
};
|
|
sessionPool = new SessionPool(opts);
|
|
await sessionPool.initialize();
|
|
await sessionPool.teardown();
|
|
|
|
entries(opts)
|
|
.filter(([key]) => key !== 'sessionOptions')
|
|
.forEach(([key, value]) => {
|
|
expect(sessionPool[key]).toEqual(value);
|
|
});
|
|
// log is appended to sessionOptions after sessionPool instantiation
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessionOptions).toEqual({ ...opts.sessionOptions, log: expect.any(Log) });
|
|
});
|
|
|
|
test('should work using SessionPool.open', async () => {
|
|
const opts = {
|
|
maxPoolSize: 3000,
|
|
|
|
sessionOptions: {
|
|
maxAgeSecs: 100,
|
|
maxUsageCount: 1,
|
|
},
|
|
|
|
persistStateKeyValueStoreId: 'TEST',
|
|
persistStateKey: 'SESSION_POOL_STATE2',
|
|
|
|
createSessionFunction: () => ({}) as never,
|
|
};
|
|
sessionPool = await SessionPool.open(opts);
|
|
await sessionPool.teardown();
|
|
|
|
entries(opts)
|
|
.filter(([key]) => key !== 'sessionOptions')
|
|
.forEach(([key, value]) => {
|
|
expect(sessionPool[key]).toEqual(value);
|
|
});
|
|
// log is appended to sessionOptions after sessionPool instantiation
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessionOptions).toEqual({ ...opts.sessionOptions, log: expect.any(Log) });
|
|
});
|
|
|
|
describe('should retrieve session', () => {
|
|
test('should retrieve session with correct shape', async () => {
|
|
sessionPool = await SessionPool.open({ sessionOptions: { maxAgeSecs: 100, maxUsageCount: 10 } });
|
|
const session = await sessionPool.getSession();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toBe(1);
|
|
expect(session.id).toBeDefined();
|
|
// @ts-expect-error Accessing private property
|
|
expect(session.maxAgeSecs).toEqual(sessionPool.sessionOptions.maxAgeSecs);
|
|
// @ts-expect-error Accessing private property
|
|
expect(session.maxUsageCount).toEqual(sessionPool.sessionOptions.maxUsageCount);
|
|
// @ts-expect-error Accessing private property
|
|
expect(session.sessionPool).toEqual(sessionPool);
|
|
});
|
|
|
|
test('should pick session when pool is full', async () => {
|
|
// @ts-expect-error private symbol
|
|
sessionPool.maxPoolSize = 2;
|
|
await sessionPool.getSession();
|
|
await sessionPool.getSession();
|
|
let isCalled = false;
|
|
// @ts-expect-error Accessing private property
|
|
const oldPick = sessionPool._pickSession;
|
|
|
|
// @ts-expect-error Overriding private property
|
|
sessionPool._pickSession = () => {
|
|
isCalled = true;
|
|
return oldPick.bind(sessionPool)();
|
|
};
|
|
|
|
await sessionPool.getSession();
|
|
|
|
expect(isCalled).toBe(true);
|
|
});
|
|
|
|
test('should delete picked session when it is unusable and create a new one', async () => {
|
|
// @ts-expect-error private symbol
|
|
sessionPool.maxPoolSize = 1;
|
|
await sessionPool.addSession();
|
|
|
|
const session = await sessionPool.getSession();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions[0].id === session.id).toBe(true);
|
|
|
|
// @ts-expect-error Overriding private property
|
|
session._errorScore += session.maxErrorScore;
|
|
await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions[0].id === session.id).toBe(false);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions).toHaveLength(1);
|
|
});
|
|
});
|
|
|
|
test('get state should work', async () => {
|
|
const url = 'https://example.com';
|
|
const cookies = [
|
|
{ name: 'cookie1', value: 'my-cookie' },
|
|
{ name: 'cookie2', value: 'your-cookie' },
|
|
];
|
|
|
|
const newSession = await sessionPool.getSession();
|
|
newSession.setCookies(cookies, url);
|
|
|
|
const state = sessionPool.getState();
|
|
expect(state).toBeInstanceOf(Object);
|
|
expect(state).toHaveProperty('usableSessionsCount');
|
|
expect(state).toHaveProperty('retiredSessionsCount');
|
|
expect(state).toHaveProperty('sessions');
|
|
});
|
|
|
|
test('should persist state and recreate it from storage', async () => {
|
|
await sessionPool.getSession();
|
|
await sessionPool.persistState();
|
|
|
|
const kvStore = await KeyValueStore.open();
|
|
const sessionPoolSaved = await kvStore.getValue<ReturnType<SessionPool['getState']>>(
|
|
// @ts-expect-error private symbol
|
|
sessionPool.persistStateKey,
|
|
);
|
|
|
|
entries(sessionPoolSaved!).forEach(([key, value]) => {
|
|
if (key !== 'sessions') {
|
|
expect(value).toEqual(sessionPool[key]);
|
|
}
|
|
});
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPoolSaved.sessions.length).toEqual(sessionPool.sessions.length);
|
|
|
|
sessionPoolSaved!.sessions.forEach((session, index) => {
|
|
entries(session).forEach(([key, value]) => {
|
|
// @ts-expect-error private symbol
|
|
if (sessionPool.sessions[index][key] instanceof Date) {
|
|
// @ts-expect-error private symbol
|
|
expect(value).toEqual((sessionPool.sessions[index][key] as Date).toISOString());
|
|
} else if (key === 'cookieJar') {
|
|
// @ts-expect-error private symbol
|
|
expect(value).toEqual(sessionPool.sessions[index][key].toJSON());
|
|
} else {
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions[index][key]).toEqual(value);
|
|
}
|
|
});
|
|
});
|
|
|
|
const loadedSessionPool = new SessionPool();
|
|
|
|
await loadedSessionPool.initialize();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions).toHaveLength(loadedSessionPool.sessions.length);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.maxPoolSize).toEqual(loadedSessionPool.maxPoolSize);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.persistStateKey).toEqual(loadedSessionPool.persistStateKey);
|
|
await sessionPool.teardown();
|
|
});
|
|
|
|
test('should create only maxPoolSize number of sessions', async () => {
|
|
// @ts-expect-error private symbol
|
|
const max = sessionPool.maxPoolSize;
|
|
for (let i = 0; i < max + 100; i++) {
|
|
await sessionPool.getSession();
|
|
}
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toEqual(sessionPool.maxPoolSize);
|
|
});
|
|
|
|
test('should create session', async () => {
|
|
// @ts-expect-error private symbol
|
|
await sessionPool._createSession();
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toBe(1);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions[0].id).toBeDefined();
|
|
});
|
|
|
|
describe('should persist state', () => {
|
|
const KV_STORE = 'SESSION-TEST';
|
|
|
|
beforeEach(async () => {
|
|
sessionPool = new SessionPool({ persistStateKeyValueStoreId: KV_STORE });
|
|
await sessionPool.initialize();
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await sessionPool.teardown();
|
|
});
|
|
|
|
test('on persist event', async () => {
|
|
await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toBe(1);
|
|
|
|
events.emit(EventType.PERSIST_STATE);
|
|
|
|
await new Promise<void>((resolve) => {
|
|
const interval = setInterval(async () => {
|
|
// @ts-expect-error private symbol
|
|
const state = await sessionPool.keyValueStore.getValue(sessionPool.persistStateKey);
|
|
if (state) {
|
|
resolve();
|
|
clearInterval(interval);
|
|
}
|
|
}, 100);
|
|
});
|
|
|
|
// @ts-expect-error private symbol
|
|
const state = await sessionPool.keyValueStore.getValue(sessionPool.persistStateKey);
|
|
|
|
expect(sessionPool.getState()).toEqual(state);
|
|
});
|
|
});
|
|
|
|
test('should remove retired sessions', async () => {
|
|
// @ts-expect-error private symbol
|
|
sessionPool.maxPoolSize = 1;
|
|
await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
const session = sessionPool.sessions[0];
|
|
// @ts-expect-error Overriding private property
|
|
session._errorScore += session.maxErrorScore;
|
|
const { id: retiredSessionId } = session;
|
|
|
|
await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.find((s) => s.id === retiredSessionId)).toEqual(undefined);
|
|
});
|
|
|
|
test('should recreate only usable sessions', async () => {
|
|
let invalidSessionsCount = 0;
|
|
for (let i = 0; i < 10; i++) {
|
|
const session = await sessionPool.getSession();
|
|
|
|
if (i % 2 === 0) {
|
|
// @ts-expect-error Overriding private property
|
|
session._errorScore += session.maxErrorScore;
|
|
invalidSessionsCount += 1;
|
|
}
|
|
}
|
|
expect(sessionPool.retiredSessionsCount).toEqual(invalidSessionsCount);
|
|
|
|
await sessionPool.persistState();
|
|
|
|
const newSessionPool = new SessionPool();
|
|
await newSessionPool.initialize();
|
|
// @ts-expect-error private symbol
|
|
expect(newSessionPool.sessions).toHaveLength(10 - invalidSessionsCount);
|
|
|
|
await newSessionPool.teardown();
|
|
});
|
|
|
|
test('should restore persisted maxUsageCount of recreated sessions', async () => {
|
|
sessionPool = await SessionPool.open({ maxPoolSize: 1, sessionOptions: { maxUsageCount: 66 } });
|
|
await sessionPool.getSession();
|
|
await sessionPool.persistState();
|
|
const loadedSessionPool = new SessionPool({ maxPoolSize: 1, sessionOptions: { maxUsageCount: 88 } });
|
|
await loadedSessionPool.initialize();
|
|
|
|
const recreatedSession = await loadedSessionPool.getSession();
|
|
|
|
expect(recreatedSession.maxUsageCount).toEqual(66);
|
|
});
|
|
|
|
test('should persist state on teardown', async () => {
|
|
const persistStateKey = 'TEST-KEY';
|
|
const persistStateKeyValueStoreId = 'TEST-VALUE-STORE';
|
|
|
|
const newSessionPool = await SessionPool.open({
|
|
maxPoolSize: 1,
|
|
persistStateKeyValueStoreId,
|
|
persistStateKey,
|
|
});
|
|
|
|
await newSessionPool.teardown();
|
|
|
|
// @ts-expect-error private symbol
|
|
const kvStore = await KeyValueStore.open(newSessionPool.persistStateKeyValueStoreId);
|
|
// @ts-expect-error private symbol
|
|
const state = await kvStore.getValue(newSessionPool.persistStateKey);
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(newSessionPool.persistStateKeyValueStoreId).toBeDefined();
|
|
// @ts-expect-error private symbol
|
|
expect(newSessionPool.persistStateKey).toBeDefined();
|
|
expect(state).toBeDefined();
|
|
expect(state).toBeInstanceOf(Object);
|
|
expect(state).toHaveProperty('usableSessionsCount');
|
|
expect(state).toHaveProperty('retiredSessionsCount');
|
|
expect(state).toHaveProperty('sessions');
|
|
});
|
|
|
|
test('should createSessionFunction work', async () => {
|
|
let isCalled;
|
|
const createSessionFunction = (sessionPool2: SessionPool) => {
|
|
isCalled = true;
|
|
expect(sessionPool2 instanceof SessionPool).toBe(true);
|
|
return new Session({ sessionPool: sessionPool2 });
|
|
};
|
|
const newSessionPool = await SessionPool.open({ createSessionFunction });
|
|
const session = await newSessionPool.getSession();
|
|
expect(isCalled).toBe(true);
|
|
expect(session.constructor.name).toBe('Session');
|
|
});
|
|
|
|
it('should remove persist state event listener', async () => {
|
|
expect(events.listenerCount(EventType.PERSIST_STATE)).toEqual(1);
|
|
await sessionPool.teardown();
|
|
expect(events.listenerCount(EventType.PERSIST_STATE)).toEqual(0);
|
|
});
|
|
|
|
test('should be able to create session with provided id', async () => {
|
|
await sessionPool.addSession({ id: 'test-session' });
|
|
// @ts-expect-error private symbol
|
|
const session = sessionPool.sessions[0];
|
|
expect(session.id).toBe('test-session');
|
|
});
|
|
|
|
test('should be able to add session instance and create new session with provided sessionOptions with addSession() ', async () => {
|
|
const session = new Session({ sessionPool, id: 'test-session-instance' });
|
|
await sessionPool.addSession(session);
|
|
|
|
await sessionPool.addSession({ id: 'test-session' });
|
|
|
|
expect(await sessionPool.getSession('test-session')).toBeDefined();
|
|
expect(await sessionPool.getSession('test-session-instance')).toBeDefined();
|
|
});
|
|
|
|
test('should not be able to add session to the pool with id already in the pool', async () => {
|
|
try {
|
|
await sessionPool.addSession({ id: 'test-session' });
|
|
await sessionPool.addSession({ id: 'test-session' });
|
|
} catch (e) {
|
|
expect((e as Error).message).toBe(
|
|
"Cannot add session with id 'test-session' as it already exists in the pool",
|
|
);
|
|
}
|
|
expect.assertions(1);
|
|
});
|
|
|
|
test('should be able to retrieve session with provided id', async () => {
|
|
await sessionPool.addSession();
|
|
await sessionPool.addSession({ id: 'test-session' });
|
|
await sessionPool.addSession({ id: 'another-test-session' });
|
|
|
|
const session = await sessionPool.getSession('test-session');
|
|
expect(session.id).toBe('test-session');
|
|
});
|
|
|
|
test('should correctly populate session array and session map', async () => {
|
|
// @ts-expect-error private symbol
|
|
sessionPool.maxPoolSize = 10;
|
|
|
|
for (let i = 0; i < 20; i++) await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toEqual(10);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessionMap.size).toEqual(10);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toEqual(sessionPool.sessionMap.size);
|
|
});
|
|
|
|
test('should correctly remove retired sessions both from array and session map', async () => {
|
|
// @ts-expect-error private symbol
|
|
sessionPool.maxPoolSize = 10;
|
|
|
|
for (let i = 0; i < 10; i++) {
|
|
await sessionPool.addSession({ id: `session_${i}` });
|
|
const session = await sessionPool.getSession(`session_${i}`);
|
|
// @ts-expect-error Overriding private property
|
|
session._errorScore += session.maxErrorScore;
|
|
}
|
|
|
|
await sessionPool.getSession();
|
|
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toEqual(1);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessionMap.size).toEqual(1);
|
|
// @ts-expect-error private symbol
|
|
expect(sessionPool.sessions.length).toEqual(sessionPool.sessionMap.size);
|
|
});
|
|
});
|