|
| 1 | +import { readFileSync } from "fs"; |
| 2 | + |
| 3 | +import { UnstructuredClient } from "../../src"; |
| 4 | +import { PartitionResponse } from "../../src/sdk/models/operations"; |
| 5 | +import { PartitionParameters, Strategy } from "../../src/sdk/models/shared"; |
| 6 | +import { describe, it, expect, vi, beforeEach} from 'vitest'; |
| 7 | + |
| 8 | +describe("FixArrayParamsHook unit tests", () => { |
| 9 | + beforeEach(() => { |
| 10 | + // Reset the mock before each test |
| 11 | + vi.resetAllMocks(); |
| 12 | + }); |
| 13 | + |
| 14 | + // Assert that array parameters are sent in the correct format |
| 15 | + // This should work with and without pdf splitting |
| 16 | + it.each([ |
| 17 | + {splitPdfPage: false}, |
| 18 | + {splitPdfPage: true}, |
| 19 | + ])( |
| 20 | + "should send extract_image_block_types in the correct format", async ({splitPdfPage}) => { |
| 21 | + const client = new UnstructuredClient({}); |
| 22 | + |
| 23 | + const file = { |
| 24 | + content: readFileSync("test/data/layout-parser-paper-fast.pdf"), |
| 25 | + fileName: "test/data/layout-parser-paper-fast.pdf", |
| 26 | + }; |
| 27 | + |
| 28 | + const requestParams: PartitionParameters = { |
| 29 | + files: file, |
| 30 | + strategy: Strategy.Fast, |
| 31 | + extractImageBlockTypes: ["a", "b", "c"], |
| 32 | + splitPdfPage: splitPdfPage, |
| 33 | + }; |
| 34 | + |
| 35 | + const fetchMock = vi.fn().mockResolvedValue( |
| 36 | + new Response( |
| 37 | + JSON.stringify([ |
| 38 | + { |
| 39 | + type: "Image", |
| 40 | + element_id: "2fe9cbfbf0ff1bd64cc4705347dbd1d6", |
| 41 | + text: "This is a test", |
| 42 | + metadata: {}, |
| 43 | + }, |
| 44 | + ]), |
| 45 | + { |
| 46 | + status: 200, |
| 47 | + headers: { "Content-Type": "application/json" }, |
| 48 | + } |
| 49 | + ) |
| 50 | + ); |
| 51 | + |
| 52 | + vi.stubGlobal("fetch", fetchMock); |
| 53 | + |
| 54 | + const res: PartitionResponse = await client.general.partition({ |
| 55 | + partitionParameters: requestParams, |
| 56 | + }); |
| 57 | + |
| 58 | + expect(fetchMock).toHaveBeenCalledTimes(1); |
| 59 | + |
| 60 | + const request = fetchMock.mock.calls[0][0]; |
| 61 | + const formData = await request.formData(); |
| 62 | + const extract_image_block_types = formData.getAll( |
| 63 | + "extract_image_block_types[]" |
| 64 | + ); |
| 65 | + |
| 66 | + expect(extract_image_block_types).toEqual(["a", "b", "c"]); |
| 67 | + }); |
| 68 | +}); |
0 commit comments