103 lines
3.9 KiB
JavaScript
103 lines
3.9 KiB
JavaScript
const fs = require('fs/promises');
|
|
const os = require('os');
|
|
const path = require('path');
|
|
const { mockClient } = require('aws-sdk-client-mock');
|
|
const { TextractClient, StartDocumentAnalysisCommand, GetDocumentAnalysisCommand } = require('@aws-sdk/client-textract');
|
|
const { LocalFsStorage } = require('../src/storage');
|
|
const { runOcr, pollDocumentAnalysis, textLayerRatio } = require('../src/ocr');
|
|
const { createLogger } = require('../src/logger');
|
|
|
|
const textract = mockClient(TextractClient);
|
|
const logger = createLogger('silent');
|
|
const config = { textLayerMinChars: 20, textLayerMinRatio: 0.8 };
|
|
|
|
describe('text layer ratio', () => {
|
|
it('counts pages above the character floor', () => {
|
|
expect(textLayerRatio([100, 0, 80, 5], 20)).toBe(0.5);
|
|
expect(textLayerRatio([], 20)).toBe(0);
|
|
});
|
|
});
|
|
|
|
describe('textract polling', () => {
|
|
let dir;
|
|
let storage;
|
|
let client;
|
|
|
|
beforeEach(async () => {
|
|
textract.reset();
|
|
dir = await fs.mkdtemp(path.join(os.tmpdir(), 'ocr-'));
|
|
storage = new LocalFsStorage(dir);
|
|
client = new TextractClient({ region: 'ap-south-1' });
|
|
});
|
|
|
|
afterEach(async () => {
|
|
await fs.rm(dir, { recursive: true, force: true });
|
|
});
|
|
|
|
it('paginates after the job succeeds', async () => {
|
|
let calls = 0;
|
|
textract.on(GetDocumentAnalysisCommand).callsFake((input) => {
|
|
calls += 1;
|
|
if (!input.NextToken && calls === 1) return { JobStatus: 'IN_PROGRESS' };
|
|
if (!input.NextToken) {
|
|
return {
|
|
JobStatus: 'SUCCEEDED',
|
|
NextToken: 'token-2',
|
|
Blocks: [{ Id: 'a', BlockType: 'LINE', Text: 'Hi', Page: 1 }],
|
|
};
|
|
}
|
|
return { JobStatus: 'SUCCEEDED', Blocks: [{ Id: 'b', BlockType: 'LINE', Text: 'There', Page: 2 }] };
|
|
});
|
|
const pages = await pollDocumentAnalysis(client, 'job-1', { sleep: async () => {} });
|
|
expect(pages).toHaveLength(2);
|
|
expect(pages[1].Blocks[0].Text).toBe('There');
|
|
});
|
|
|
|
it('throws on FAILED and PARTIAL_SUCCESS', async () => {
|
|
textract.on(GetDocumentAnalysisCommand).resolvesOnce({ JobStatus: 'FAILED', StatusMessage: 'bad scan' });
|
|
await expect(pollDocumentAnalysis(client, 'job-2', { sleep: async () => {} })).rejects.toThrow(/FAILED: bad scan/);
|
|
textract.reset();
|
|
textract.on(GetDocumentAnalysisCommand).resolvesOnce({ JobStatus: 'PARTIAL_SUCCESS', StatusMessage: 'cut off' });
|
|
await expect(pollDocumentAnalysis(client, 'job-3', { sleep: async () => {} })).rejects.toThrow(/PARTIAL_SUCCESS: cut off/);
|
|
});
|
|
|
|
it('skips Textract when raw output already exists', async () => {
|
|
await storage.putBytes('books/9/source.pdf', Buffer.from('%PDF'));
|
|
await storage.putJson('books/9/textract/raw_0001.json', { Blocks: [] });
|
|
textract.on(StartDocumentAnalysisCommand).resolves({ JobId: 'should-not-run' });
|
|
const result = await runOcr({
|
|
bookId: 9,
|
|
force: false,
|
|
storage,
|
|
textract: client,
|
|
bucket: 'bucket',
|
|
config,
|
|
logger,
|
|
sleep: async () => {},
|
|
});
|
|
expect(result).toMatchObject({ skipped: true, method: 'cached' });
|
|
expect(textract.commandCalls(StartDocumentAnalysisCommand)).toHaveLength(0);
|
|
});
|
|
|
|
it('loads a textract json file and does not start a job', async () => {
|
|
const pdf = path.join(dir, 'book.pdf');
|
|
const raw = path.join(dir, 'given.json');
|
|
await fs.writeFile(pdf, '%PDF');
|
|
await fs.writeFile(raw, JSON.stringify({ JobStatus: 'SUCCEEDED', Blocks: [{ Id: 'a', BlockType: 'LINE', Text: 'Saved', Page: 1 }] }));
|
|
const result = await runOcr({
|
|
bookId: 4,
|
|
pdf,
|
|
textractJson: raw,
|
|
storage,
|
|
textract: client,
|
|
bucket: 'bucket',
|
|
config,
|
|
logger,
|
|
});
|
|
expect(result.method).toBe('textract-json');
|
|
const stored = await storage.getJson('books/4/textract/raw_0001.json');
|
|
expect(stored.Blocks[0].Text).toBe('Saved');
|
|
expect(textract.commandCalls(StartDocumentAnalysisCommand)).toHaveLength(0);
|
|
});
|
|
});
|