mynode/packages/lib/services/ocr/OcrService.test.ts

import { createNoteAndResource, newOcrService, ocrSampleDir, resourceFetcher, setupDatabaseAndSynchronizer, supportDir, switchClient, synchronizerStart, withWarningSilenced } from '../../testing/test-utils';
import { supportedMimeTypes } from './OcrService';
import Resource from '../../models/Resource';
import { ResourceEntity, ResourceOcrDriverId, ResourceOcrStatus } from '../database/types';
import { msleep } from '@joplin/utils/time';
import Logger from '@joplin/utils/Logger';
import Setting from '../../models/Setting';
import createAccessiblePdf from './utils/createAccessiblePdf';
import { PdfOcrDetails } from './utils/types';
import * as fs from 'fs-extra';
import shim from '../../shim';
import { copyFile } from 'fs/promises';
import { join } from 'path';

const highConfidenceTestOcrResource = {
	expectedText: 'This is a lot of 12 point text to test the\n' +
		'ocr code and see if it works on all types\n' +
		'of file format.\n' +
		'The quick brown dog jumped over the\n' +
		'lazy fox. The quick brown dog jumped\n' +
		'over the lazy fox. The quick brown dog\n' +
		'jumped over the lazy fox. The quick\n' +
		'brown dog jumped over the lazy fox.',
	path: `${ocrSampleDir}/testocr.png`,
};

describe('OcrService', () => {

	jest.retryTimes(2);

	beforeEach(async () => {
		await setupDatabaseAndSynchronizer(1);
		await setupDatabaseAndSynchronizer(2);
		await switchClient(1);
	});

	it('should process resources', async () => {
		const { resource: resource1 } = await createNoteAndResource({ path: highConfidenceTestOcrResource.path });
		const { resource: resource2 } = await createNoteAndResource({ path: `${supportDir}/photo.jpg` });
		const { resource: resource3 } = await createNoteAndResource({ path: `${ocrSampleDir}/with_bullets.png` });

		// Wait to make sure that updated_time is updated
		await msleep(1);

		expect(await Resource.needOcrCount(supportedMimeTypes)).toBe(3);

		const service = newOcrService();
		await service.processResources();

		const processedResource1: ResourceEntity = await Resource.load(resource1.id);
		expect(processedResource1.ocr_text).toBe(highConfidenceTestOcrResource.expectedText);
		expect(processedResource1.ocr_status).toBe(ResourceOcrStatus.Done);
		expect(processedResource1.ocr_error).toBe('');

		const details = Resource.unserializeOcrDetails(processedResource1.ocr_details);
		const lines = details.map(l => l.words.map(w => w.t).join(' ')).join('\n');
		expect(lines).toBe(highConfidenceTestOcrResource.expectedText);
		expect(details[0].words[0].t).toBe('This');
		expect(details[0].words[0]).toEqual({ 't': 'This', 'bb': [36, 96, 92, 116], 'bl': [36, 96, 116, 116] });

		// Also check that the resource blob has not been updated
		expect(processedResource1.blob_updated_time).toBe(resource1.blob_updated_time);
		expect(processedResource1.updated_time).toBeGreaterThan(resource1.updated_time);

		const processedResource2: ResourceEntity = await Resource.load(resource2.id);
		expect(processedResource2.ocr_text).toBe('');
		expect(processedResource2.ocr_status).toBe(ResourceOcrStatus.Done);
		expect(processedResource2.ocr_error).toBe('');

		const processedResource3: ResourceEntity = await Resource.load(resource3.id);
		expect(processedResource3.ocr_text).toBe('Declaration\n' +
			'| declare that:\n' +
			'® | will arrive in the UK within the next 48 hours\n' +
			'® | understand | have to provide proof of a negative COVID 19 test prior to departure to the UK (unless\n' +
			'exempt)\n' +
			'® | have provided my seat number, if relevant\n' +
			'® The information | have entered in this form is correct\n' +
			'® | understand it could be a criminal offence to provide false details and | may be prosecuted\n' +
			'If any of your information changes once you have submitted your details, such as travel details, seat number, or\n' +
			'contact information, you must complete a new form.\n' +
			'| confirm that | understand and agree with the above declarations.',
		);
		expect(processedResource3.ocr_status).toBe(ResourceOcrStatus.Done);
		expect(processedResource3.ocr_error).toBe('');

		// Also check that the resource blob has not been updated
		expect(processedResource2.blob_updated_time).toBe(resource2.blob_updated_time);
		expect(processedResource2.updated_time).toBeGreaterThan(resource2.updated_time);

		await service.dispose();

		// On CI these tests can randomly throw the error "Exceeded timeout of
		// 90000 ms for a test.". So for now increase the timeout and if that's
		// not sufficient it means the test is simply stuck, and we should use
		// `jest.retryTimes(2)`
	}, 60000 * 5);

	test.each([
		// Use embedded text (skip OCR)
		['dummy.pdf', 'Dummy PDF file'],
		['multi_page__embedded_text.pdf', 'This is a test.\nTesting...\nThis PDF has 3 pages.\nThis is page 3.'],
	])('should process PDF resources', async (samplePath: string, expectedText: string) => {
		const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/${samplePath}` });

		const service = newOcrService();

		await service.processResources();

		const processedResource: ResourceEntity = await Resource.load(resource.id);
		expect(processedResource.ocr_text).toBe(expectedText);
		expect(processedResource.ocr_status).toBe(ResourceOcrStatus.Done);
		expect(processedResource.ocr_error).toBe('');

		await service.dispose();
	});

	it('should handle case where resource blob has not yet been downloaded', async () => {
		await createNoteAndResource({ path: `${ocrSampleDir}/dummy.pdf` });

		await synchronizerStart();

		await switchClient(2);

		await synchronizerStart();

		await msleep(1);

		const service = newOcrService();

		await service.processResources();

		{
			const resource: ResourceEntity = (await Resource.all())[0];
			expect(resource.ocr_text).toBe('');
			expect(resource.ocr_error).toBe('');
			expect(resource.ocr_status).toBe(ResourceOcrStatus.Todo);
		}

		await resourceFetcher().startAndWait();

		await service.processResources();

		{
			const resource: ResourceEntity = (await Resource.all())[0];
			expect(resource.ocr_text).toBe('Dummy PDF file');
			expect(resource.ocr_error).toBe('');
			expect(resource.ocr_status).toBe(ResourceOcrStatus.Done);
		}

		await service.dispose();
	});

	it('should handle case where resource blob cannot be downloaded', async () => {
		await createNoteAndResource({ path: `${ocrSampleDir}/dummy.pdf` });

		await synchronizerStart();

		await switchClient(2);

		await synchronizerStart();

		const resource: ResourceEntity = (await Resource.all())[0];

		// ----------------------------------------------------------------
		// Fetch status is an error so OCR status will be an error too
		// ----------------------------------------------------------------

		await Resource.setLocalState(resource.id, {
			resource_id: resource.id,
			fetch_status: Resource.FETCH_STATUS_ERROR,
			fetch_error: 'cannot be downloaded',
		});

		const service = newOcrService();

		// The service will print a warning so we disable it in tests
		Logger.globalLogger.enabled = false;
		await service.processResources();
		Logger.globalLogger.enabled = true;

		{
			const resource: ResourceEntity = (await Resource.all())[0];
			expect(resource.ocr_text).toBe('');
			expect(resource.ocr_error).toContain('Cannot process resource');
			expect(resource.ocr_error).toContain('cannot be downloaded');
			expect(resource.ocr_status).toBe(ResourceOcrStatus.Error);
		}

		// ----------------------------------------------------------------
		// After the fetch status is reset and the resource downloaded, it
		// should also retry OCR and succeed.
		// ----------------------------------------------------------------

		await Resource.resetFetchErrorStatus(resource.id);

		await resourceFetcher().startAndWait();

		await service.processResources();

		{
			const resource: ResourceEntity = (await Resource.all())[0];
			expect(resource.ocr_text).toBe('Dummy PDF file');
			expect(resource.ocr_error).toBe('');
			expect(resource.ocr_status).toBe(ResourceOcrStatus.Done);
		}

		await service.dispose();
	});

	it('should handle conflicts if two clients process the same resource then sync', async () => {
		await createNoteAndResource({ path: `${ocrSampleDir}/dummy.pdf` });

		const service1 = newOcrService();
		await synchronizerStart();
		await service1.processResources();

		await switchClient(2);

		await synchronizerStart();
		await msleep(1);
		await resourceFetcher().startAndWait();
		const service2 = newOcrService();
		await service2.processResources();
		await synchronizerStart();
		const expectedResourceUpdatedTime = (await Resource.all())[0].updated_time;

		await switchClient(1);

		await synchronizerStart();

		// A conflict happened during sync, but it is resolved by keeping the
		// remote version.

		expect((await Resource.all()).length).toBe(1);

		{
			const resource: ResourceEntity = (await Resource.all())[0];
			expect(resource.ocr_text).toBe('Dummy PDF file');
			expect(resource.ocr_error).toBe('');
			expect(resource.ocr_status).toBe(ResourceOcrStatus.Done);
			expect(resource.updated_time).toBe(expectedResourceUpdatedTime);
		}

		await service1.dispose();
		await service2.dispose();
	});

	// Use this to quickly test with specific images:

	// it('should process resources 2', async () => {
	// 	await createNoteAndResource({ path: `${require('os').homedir()}/Desktop/AllClients.png` });

	// 	const service = newOcrService();
	// 	await service.processResources();

	// 	console.info(await Resource.all());

	// 	await service.dispose();
	// });

	it('should generate text even on cases of lower confidence', async () => {
		const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/low_confidence_testing.png` });

		const service = newOcrService();
		await service.processResources();

		const processedResource: ResourceEntity = await Resource.load(resource.id);
		expect(processedResource.ocr_text.includes('1.')).toBe(true);
		// cSpell:disable
		expect(processedResource.ocr_text.includes('eback Mountain (2005)')).toBe(true);
		// cSpell:enable

		expect(processedResource.ocr_text.includes('2.')).toBe(true);
		expect(processedResource.ocr_text.includes('Havoc (2005)')).toBe(true);

		expect(processedResource.ocr_text.includes('3.')).toBe(true);
		expect(processedResource.ocr_text.includes('Love & Other Drugs (2010)')).toBe(true);

		expect(processedResource.ocr_text.includes('4.')).toBe(true);
		expect(processedResource.ocr_text.includes('The Last Thing He Wanted (2020)')).toBe(true);

		await service.dispose();
	});

	it('should skip resources with an invalid ocr_driver_id', async () => {
		const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/dummy.pdf` });

		await Resource.save({
			...resource,
			ocr_driver_id: -123456, // An invalid ID
		});

		const service = newOcrService();

		// Should not loop forever
		await service.processResources();

		const processedResource: ResourceEntity = await Resource.load(resource.id);
		expect(processedResource.ocr_text).toBe('');

		await service.dispose();
	});

	it('should process resources with ocr_driver_id 0 as printed text', async () => {
		const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/multi_page__embedded_text.pdf` });
		await Resource.save({
			...resource,
			ocr_driver_id: 0,
		});

		const service = newOcrService();
		await service.processResources();

		const processedResource: ResourceEntity = await Resource.load(resource.id);
		expect(processedResource.ocr_text).toBe('This is a test.\nTesting...\nThis PDF has 3 pages.\nThis is page 3.');
		expect(processedResource.ocr_status).toBe(ResourceOcrStatus.Done);
		expect(processedResource.ocr_error).toBe('');

		await service.dispose();
	});

	it('should skip HTR processing when the relevant setting is disabled', async () => {
		const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/multi_page__embedded_text.pdf` });
		Setting.setValue('ocr.handwrittenTextDriverEnabled', false);

		await Resource.save({
			...resource,
			ocr_driver_id: ResourceOcrDriverId.HandwrittenText,
			title: 'Test',
		});

		const service = newOcrService();
		await service.processResources();

		// Should not process HandwrittenText results
		const processedResource: ResourceEntity = await Resource.load(resource.id);
		expect(processedResource).toMatchObject({
			ocr_text: '',
			title: 'Test',
			ocr_status: ResourceOcrStatus.Todo,
			ocr_error: '',
		});

		await service.dispose();
	});

	it('should throw error for unsupported OCR details version', async () => {
		const ocrDetails = {
			version: 999,
			pages: [] as PdfOcrDetails['pages'],
		};

		await expect(createAccessiblePdf([], JSON.stringify(ocrDetails)))
			.rejects.toThrow('Unsupported PDF OCR details version: 999');
	});

	it('should throw error for page count mismatch', async () => {
		const ocrDetails: PdfOcrDetails = {
			version: 1,
			pages: [{ lines: [] }],
		};

		await expect(createAccessiblePdf([], JSON.stringify(ocrDetails)))
			.rejects.toThrow('Page count mismatch: 0 images vs 1 OCR pages');
	});

	it('should create a multi-page PDF', async () => {
		const jpegBuffer = await fs.readFile(`${supportDir}/photo.jpg`);

		const ocrDetails: PdfOcrDetails = {
			version: 1,
			pages: [
				{
					lines: [{ words: [{ t: 'Page1', bb: [10, 60, 10, 30] }] }],
				},
				{
					lines: [{ words: [{ t: 'Page2', bb: [10, 60, 10, 30] }] }],
				},
			],
		};

		const pdfBytes = await createAccessiblePdf(
			[
				{ buffer: jpegBuffer, width: 200, height: 200 },
				{ buffer: jpegBuffer, width: 200, height: 200 },
			],
			JSON.stringify(ocrDetails),
		);

		expect(pdfBytes).toBeInstanceOf(Uint8Array);
		const pdfContent = new TextDecoder().decode(pdfBytes);
		expect(pdfContent.startsWith('%PDF-')).toBe(true);
		expect(pdfContent).toContain('%%EOF');

		// Multi-page PDF should be roughly twice the size of single page
		// (minus some overhead for shared resources like fonts)
		expect(pdfBytes.length).toBeGreaterThan(jpegBuffer.length * 1.5);
	});

	it('should truncate long OCR text', async () => {
		const service = newOcrService();
		service.testing_setOcrMaxSize(10);

		const { resource: resource1 } = await createNoteAndResource({ path: `${ocrSampleDir}/multi_page__embedded_text.pdf` });
		const { resource: resource2 } = await createNoteAndResource({ path: `${ocrSampleDir}/testocr.png` });
		await msleep(1);
		expect(await Resource.needOcrCount(supportedMimeTypes)).toBe(2);

		await service.processResources();

		expect(await Resource.load(resource1.id)).toMatchObject({
			ocr_text: 'This is...',
		});
		// PNG resources are unlikely to go over the actual ocr_text length limit, but the logic should
		// support them anyway:
		expect(await Resource.load(resource2.id)).toMatchObject({
			ocr_text: 'This is...',
		});
	});

	it('should truncate long OCR text in accessible-mode PDFs', async () => {
		Setting.setValue('ocr.pdfMode', 'accessible');

		// Converting PDFs to images requires the <canvas> element, which isn't available in a testing
		// environment.
		const pdfToImages = jest.spyOn(shim, 'pdfToImages');
		pdfToImages.mockImplementation(async (_pdfPath, outputDirectory) => {
			const paths = [];
			for (let i = 0; i < 5; i++) {
				const outputPath = join(outputDirectory, `page-${i}.png`);
				await copyFile(`${ocrSampleDir}/testocr.png`, outputPath);
				paths.push(outputPath);
			}
			return paths;
		});

		const service = newOcrService();
		const maxSize = highConfidenceTestOcrResource.expectedText.length * 2 + '.\n...'.length;
		service.testing_setOcrMaxSize(maxSize);

		try {
			const { resource } = await createNoteAndResource({ path: `${ocrSampleDir}/multi_page__embedded_text.pdf` });
			await msleep(1);
			expect(await Resource.needOcrCount(supportedMimeTypes)).toBe(1);

			await withWarningSilenced(/Recognize: Truncated/, async () => {
				await service.processResources();
			});

			const reloaded = await Resource.load(resource.id);
			expect(reloaded).toMatchObject({
				ocr_text: [
					highConfidenceTestOcrResource.expectedText,
					highConfidenceTestOcrResource.expectedText,
					'...',
				].join('\n'),
			});

			// Should **not** truncate ocr_details
			const parsedOcrDetails: PdfOcrDetails = JSON.parse(reloaded.ocr_details);
			expect(parsedOcrDetails.pages).toHaveLength(5);
		} finally {
			pdfToImages.mockRestore();
		}
	});

});