|
| 1 | +/** |
| 2 | + * imageModelIntegrity IO wrappers — resolveImageModelDir / validateImageModelDir / |
| 3 | + * ensureImageExtractionComplete. The pure rule is covered in imageModelIntegrity.test; |
| 4 | + * this drives the RNFS + unzip boundaries (mocked) to cover dir resolution, the coreml |
| 5 | + * short-circuit, unreadable dirs, and the re-unzip-once-then-throw gate. |
| 6 | + */ |
| 7 | +const existing = new Set<string>(); |
| 8 | +const dirListings: Record<string, Array<{ name: string; size: number; isFile: boolean }>> = {}; |
| 9 | + |
| 10 | +const mockExists = jest.fn(async (p: string) => existing.has(p)); |
| 11 | +const mockReadDir = jest.fn(async (p: string) => { |
| 12 | + if (!(p in dirListings)) throw new Error(`ENOENT ${p}`); |
| 13 | + return dirListings[p].map(e => ({ |
| 14 | + name: e.name, |
| 15 | + size: e.size, |
| 16 | + path: `${p}/${e.name}`, |
| 17 | + isFile: () => e.isFile, |
| 18 | + isDirectory: () => !e.isFile, |
| 19 | + })); |
| 20 | +}); |
| 21 | +jest.mock('react-native-fs', () => ({ exists: (p: string) => mockExists(p), readDir: (p: string) => mockReadDir(p) })); |
| 22 | + |
| 23 | +const mockUnzip = jest.fn(async (_source?: string, _target?: string) => '/done'); |
| 24 | +jest.mock('react-native-zip-archive', () => ({ unzip: (source: string, target: string) => mockUnzip(source, target) })); |
| 25 | + |
| 26 | +jest.mock('../../../src/utils/logger', () => ({ __esModule: true, default: { log: jest.fn(), warn: jest.fn(), error: jest.fn() } })); |
| 27 | + |
| 28 | +import { |
| 29 | + resolveImageModelDir, validateImageModelDir, ensureImageExtractionComplete, |
| 30 | +} from '../../../src/utils/imageModelIntegrity'; |
| 31 | +import { ImageModelIncompleteError } from '../../../src/services/modelLoadErrors'; |
| 32 | + |
| 33 | +const COMPLETE_MNN = [ |
| 34 | + { name: 'unet.mnn', size: 1, isFile: true }, { name: 'unet.mnn.weight', size: 1, isFile: true }, |
| 35 | + { name: 'clip_v2.mnn', size: 1, isFile: true }, { name: 'clip_v2.mnn.weight', size: 1, isFile: true }, |
| 36 | + { name: 'vae_decoder.mnn', size: 1, isFile: true }, { name: 'vae_decoder.mnn.weight', size: 1, isFile: true }, |
| 37 | + { name: 'pos_emb.bin', size: 1, isFile: true }, { name: 'token_emb.bin', size: 1, isFile: true }, |
| 38 | + { name: 'tokenizer.json', size: 1, isFile: true }, |
| 39 | +]; |
| 40 | + |
| 41 | +beforeEach(() => { |
| 42 | + jest.clearAllMocks(); |
| 43 | + existing.clear(); |
| 44 | + for (const k of Object.keys(dirListings)) delete dirListings[k]; |
| 45 | + mockUnzip.mockResolvedValue('/done'); |
| 46 | +}); |
| 47 | + |
| 48 | +describe('resolveImageModelDir', () => { |
| 49 | + it('returns modelPath itself when the unet marker is directly inside', async () => { |
| 50 | + existing.add('/m/unet.mnn'); |
| 51 | + expect(await resolveImageModelDir('/m', 'mnn')).toBe('/m'); |
| 52 | + }); |
| 53 | + |
| 54 | + it('finds the marker one subdir down (e.g. AbsoluteReality/)', async () => { |
| 55 | + existing.add('/m/AbsoluteReality/unet.mnn'); |
| 56 | + dirListings['/m'] = [{ name: 'AbsoluteReality', size: 0, isFile: false }]; |
| 57 | + expect(await resolveImageModelDir('/m', 'mnn')).toBe('/m/AbsoluteReality'); |
| 58 | + }); |
| 59 | + |
| 60 | + it('finds the marker two subdirs down (qnn output_512/qnn_models_min)', async () => { |
| 61 | + existing.add('/m/output_512/qnn_models_min/unet.bin'); |
| 62 | + dirListings['/m'] = [{ name: 'output_512', size: 0, isFile: false }]; |
| 63 | + dirListings['/m/output_512'] = [{ name: 'qnn_models_min', size: 0, isFile: false }]; |
| 64 | + expect(await resolveImageModelDir('/m', 'qnn')).toBe('/m/output_512/qnn_models_min'); |
| 65 | + }); |
| 66 | + |
| 67 | + it('returns null when no marker is found and skips unreadable subdirs', async () => { |
| 68 | + dirListings['/m'] = [{ name: 'empty', size: 0, isFile: false }]; // /m/empty is unreadable (no listing) |
| 69 | + expect(await resolveImageModelDir('/m', 'mnn')).toBeNull(); |
| 70 | + }); |
| 71 | + |
| 72 | + it('skips top-level files and non-marker sub-files while descending two levels', async () => { |
| 73 | + existing.add('/m/sub/deep/unet.mnn'); |
| 74 | + dirListings['/m'] = [ |
| 75 | + { name: 'readme.txt', size: 5, isFile: true }, // top-level FILE → item.isDirectory() false |
| 76 | + { name: 'sub', size: 0, isFile: false }, |
| 77 | + ]; |
| 78 | + dirListings['/m/sub'] = [ |
| 79 | + { name: 'note.bin', size: 5, isFile: true }, // sub FILE → s.isDirectory() false |
| 80 | + { name: 'deep', size: 0, isFile: false }, |
| 81 | + ]; |
| 82 | + expect(await resolveImageModelDir('/m', 'mnn')).toBe('/m/sub/deep'); |
| 83 | + }); |
| 84 | + |
| 85 | + it('returns null when modelPath itself is unreadable', async () => { |
| 86 | + expect(await resolveImageModelDir('/nope', 'mnn')).toBeNull(); |
| 87 | + }); |
| 88 | + |
| 89 | + it('treats an exists() failure as no-marker (hasMarker swallows the error)', async () => { |
| 90 | + mockExists.mockRejectedValueOnce(new Error('io error')); // the top-level marker probe throws |
| 91 | + expect(await resolveImageModelDir('/m', 'mnn')).toBeNull(); // then /m is unreadable → null |
| 92 | + }); |
| 93 | +}); |
| 94 | + |
| 95 | +describe('validateImageModelDir', () => { |
| 96 | + it('coreml short-circuits to complete', async () => { |
| 97 | + expect(await validateImageModelDir('/m', 'coreml')).toEqual({ complete: true, missing: [] }); |
| 98 | + }); |
| 99 | + |
| 100 | + it('reports the unet as missing when no model dir resolves', async () => { |
| 101 | + const res = await validateImageModelDir('/m', 'mnn'); |
| 102 | + expect(res.complete).toBe(false); |
| 103 | + expect(res.missing).toContain('unet.mnn'); |
| 104 | + }); |
| 105 | + |
| 106 | + it('names unet.bin (not unet.mnn) when a qnn model dir does not resolve', async () => { |
| 107 | + const res = await validateImageModelDir('/m', 'qnn'); |
| 108 | + expect(res.missing).toContain('unet.bin'); |
| 109 | + }); |
| 110 | + |
| 111 | + it('treats a zero/NaN file size as 0 (missing) when mapping the listing', async () => { |
| 112 | + existing.add('/m/unet.mnn'); |
| 113 | + // token_emb.bin present but zero-byte → Number(size)||0 = 0 → counted missing. |
| 114 | + dirListings['/m'] = COMPLETE_MNN.map(f => (f.name === 'token_emb.bin' ? { ...f, size: 0 } : f)); |
| 115 | + const res = await validateImageModelDir('/m', 'mnn'); |
| 116 | + expect(res.complete).toBe(false); |
| 117 | + expect(res.missing).toContain('token_emb.bin'); |
| 118 | + }); |
| 119 | + |
| 120 | + it('passes a complete extraction', async () => { |
| 121 | + existing.add('/m/unet.mnn'); |
| 122 | + dirListings['/m'] = COMPLETE_MNN; |
| 123 | + expect(await validateImageModelDir('/m', 'mnn')).toEqual({ complete: true, missing: [] }); |
| 124 | + }); |
| 125 | + |
| 126 | + it('flags the missing files of a partial extraction', async () => { |
| 127 | + existing.add('/m/unet.mnn'); |
| 128 | + dirListings['/m'] = COMPLETE_MNN.filter(f => f.name !== 'pos_emb.bin' && f.name !== 'clip_v2.mnn.weight'); |
| 129 | + const res = await validateImageModelDir('/m', 'mnn'); |
| 130 | + expect(res.complete).toBe(false); |
| 131 | + expect(res.missing).toEqual(expect.arrayContaining(['pos_emb.bin', 'clip_v2.mnn.weight'])); |
| 132 | + }); |
| 133 | + |
| 134 | + it('reports unreadable when the resolved dir cannot be listed', async () => { |
| 135 | + existing.add('/m/unet.mnn'); |
| 136 | + // marker exists so it resolves to /m, but /m has no listing → readDir throws. |
| 137 | + const res = await validateImageModelDir('/m', 'mnn'); |
| 138 | + expect(res.complete).toBe(false); |
| 139 | + expect(res.missing).toContain('<unreadable model dir>'); |
| 140 | + }); |
| 141 | +}); |
| 142 | + |
| 143 | +describe('ensureImageExtractionComplete', () => { |
| 144 | + const opts = (backend?: string) => ({ backend: backend as any, modelDir: '/m', zipPath: '/m.zip', modelId: 'id' }); |
| 145 | + |
| 146 | + it('is a no-op for coreml / unknown backends', async () => { |
| 147 | + await expect(ensureImageExtractionComplete(opts('coreml'))).resolves.toBeUndefined(); |
| 148 | + await expect(ensureImageExtractionComplete(opts(undefined))).resolves.toBeUndefined(); |
| 149 | + expect(mockUnzip).not.toHaveBeenCalled(); |
| 150 | + }); |
| 151 | + |
| 152 | + it('passes without re-unzip when already complete', async () => { |
| 153 | + existing.add('/m/unet.mnn'); |
| 154 | + dirListings['/m'] = COMPLETE_MNN; |
| 155 | + await expect(ensureImageExtractionComplete(opts('mnn'))).resolves.toBeUndefined(); |
| 156 | + expect(mockUnzip).not.toHaveBeenCalled(); |
| 157 | + }); |
| 158 | + |
| 159 | + it('re-unzips ONCE and passes when the retry completes the model', async () => { |
| 160 | + existing.add('/m/unet.mnn'); |
| 161 | + dirListings['/m'] = COMPLETE_MNN.filter(f => f.name !== 'pos_emb.bin'); // incomplete first |
| 162 | + mockUnzip.mockImplementation(async () => { dirListings['/m'] = COMPLETE_MNN; return '/done'; }); // repaired on re-unzip |
| 163 | + await expect(ensureImageExtractionComplete(opts('mnn'))).resolves.toBeUndefined(); |
| 164 | + expect(mockUnzip).toHaveBeenCalledTimes(1); |
| 165 | + }); |
| 166 | + |
| 167 | + it('throws ImageModelIncompleteError when still incomplete after the re-unzip', async () => { |
| 168 | + existing.add('/m/unet.mnn'); |
| 169 | + dirListings['/m'] = COMPLETE_MNN.filter(f => f.name !== 'pos_emb.bin'); |
| 170 | + // re-unzip is a no-op (still missing pos_emb.bin) |
| 171 | + await expect(ensureImageExtractionComplete(opts('mnn'))).rejects.toBeInstanceOf(ImageModelIncompleteError); |
| 172 | + expect(mockUnzip).toHaveBeenCalledTimes(1); |
| 173 | + }); |
| 174 | +}); |
0 commit comments