Cleanup work + stream endpoints to reduce memory usage (#6106)

This commit is contained in:
Anthony Stirling
2026-04-15 15:34:17 +01:00
committed by GitHub
parent 702f4e5c2c
commit cc5a0b8def
116 changed files with 3005 additions and 1512 deletions
@@ -9,6 +9,7 @@ import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mockStatic;
import static org.mockito.Mockito.when;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
@@ -29,6 +30,7 @@ import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.mock.web.MockMultipartFile;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;
import io.github.pixee.security.ZipSecurity;
@@ -59,6 +61,19 @@ class PDFToFileTest {
.thenAnswer(
invocation ->
Files.createTempFile("test", invocation.getArgument(0)).toFile());
lenient()
.when(mockTempFileManager.createManagedTempFile(anyString()))
.thenAnswer(
invocation -> {
File f =
Files.createTempFile("test", invocation.<String>getArgument(0))
.toFile();
TempFile tf = org.mockito.Mockito.mock(TempFile.class);
lenient().when(tf.getFile()).thenReturn(f);
lenient().when(tf.getPath()).thenReturn(f.toPath());
lenient().when(tf.getAbsolutePath()).thenReturn(f.getAbsolutePath());
return tf;
});
lenient()
.when(mockTempFileManager.createTempDirectory())
.thenAnswer(invocation -> Files.createTempDirectory("test"));
@@ -68,6 +83,12 @@ class PDFToFileTest {
pdfToFile = new PDFToFile(mockTempFileManager, mockRuntimePathConfig);
}
private static byte[] drain(ResponseEntity<StreamingResponseBody> response) throws IOException {
ByteArrayOutputStream baos = new ByteArrayOutputStream();
response.getBody().writeTo(baos);
return baos.toByteArray();
}
@Test
void testProcessPdfToMarkdown_InvalidContentType() throws IOException, InterruptedException {
// Prepare
@@ -79,7 +100,7 @@ class PDFToFileTest {
"This is not a PDF".getBytes());
// Execute
ResponseEntity<byte[]> response = pdfToFile.processPdfToMarkdown(nonPdfFile);
ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToMarkdown(nonPdfFile);
// Verify
assertEquals(HttpStatus.BAD_REQUEST, response.getStatusCode());
@@ -96,7 +117,7 @@ class PDFToFileTest {
"This is not a PDF".getBytes());
// Execute
ResponseEntity<byte[]> response = pdfToFile.processPdfToHtml(nonPdfFile);
ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToHtml(nonPdfFile);
// Verify
assertEquals(HttpStatus.BAD_REQUEST, response.getStatusCode());
@@ -114,7 +135,7 @@ class PDFToFileTest {
"This is not a PDF".getBytes());
// Execute
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(nonPdfFile, "docx", "draw_pdf_import");
// Verify
@@ -133,7 +154,7 @@ class PDFToFileTest {
"Fake PDF content".getBytes());
// Execute with invalid format
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(pdfFile, "invalid_format", "draw_pdf_import");
// Verify
@@ -184,12 +205,14 @@ class PDFToFileTest {
});
// Execute the method
ResponseEntity<byte[]> response = pdfToFile.processPdfToMarkdown(pdfFile);
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToMarkdown(pdfFile);
// Verify - should now return a ZIP file instead of plain markdown
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition indicates a ZIP file
assertTrue(
@@ -201,7 +224,7 @@ class PDFToFileTest {
// Verify the content by unzipping it
try (ZipInputStream zipStream =
ZipSecurity.createHardenedInputStream(
new java.io.ByteArrayInputStream(response.getBody()))) {
new java.io.ByteArrayInputStream(bodyBytes))) {
ZipEntry entry;
boolean foundMdFile = false;
boolean foundImageInFolder = false;
@@ -275,12 +298,14 @@ class PDFToFileTest {
});
// Execute the method
ResponseEntity<byte[]> response = pdfToFile.processPdfToMarkdown(pdfFile);
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToMarkdown(pdfFile);
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition indicates a zip file
assertTrue(
@@ -292,7 +317,7 @@ class PDFToFileTest {
// Verify the content by unzipping it
try (ZipInputStream zipStream =
ZipSecurity.createHardenedInputStream(
new java.io.ByteArrayInputStream(response.getBody()))) {
new java.io.ByteArrayInputStream(bodyBytes))) {
ZipEntry entry;
boolean foundMdFiles = false;
boolean foundImage = false;
@@ -352,12 +377,13 @@ class PDFToFileTest {
});
// Execute the method
ResponseEntity<byte[]> response = pdfToFile.processPdfToHtml(pdfFile);
ResponseEntity<StreamingResponseBody> response = pdfToFile.processPdfToHtml(pdfFile);
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition indicates a zip file
assertTrue(
@@ -369,7 +395,7 @@ class PDFToFileTest {
// Verify the content by unzipping it
try (ZipInputStream zipStream =
ZipSecurity.createHardenedInputStream(
new java.io.ByteArrayInputStream(response.getBody()))) {
new java.io.ByteArrayInputStream(bodyBytes))) {
ZipEntry entry;
boolean foundMainHtml = false;
boolean foundIndexHtml = false;
@@ -437,13 +463,14 @@ class PDFToFileTest {
});
// Execute the method with docx format
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(pdfFile, "docx", "draw_pdf_import");
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition has correct filename
assertTrue(
@@ -508,13 +535,14 @@ class PDFToFileTest {
});
// Execute the method with ODP format
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(pdfFile, "odp", "draw_pdf_import");
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition for zip file
assertTrue(
@@ -526,7 +554,7 @@ class PDFToFileTest {
// Verify the content by unzipping it
try (ZipInputStream zipStream =
ZipSecurity.createHardenedInputStream(
new java.io.ByteArrayInputStream(response.getBody()))) {
new java.io.ByteArrayInputStream(bodyBytes))) {
ZipEntry entry;
boolean foundMainFile = false;
boolean foundMediaFiles = false;
@@ -592,13 +620,14 @@ class PDFToFileTest {
});
// Execute the method with text format
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(pdfFile, "txt:Text", "draw_pdf_import");
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition has txt extension
assertTrue(
@@ -650,13 +679,14 @@ class PDFToFileTest {
});
// Execute the method
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFile.processPdfToOfficeFormat(pdfFile, "docx", "draw_pdf_import");
// Verify
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
// Verify content disposition contains output.docx
assertTrue(
@@ -696,12 +726,13 @@ class PDFToFileTest {
return mockExecutorResult;
});
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFileWithUno.processPdfToOfficeFormat(pdfFile, "docx", "writer_pdf_import");
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
assertTrue(
response.getHeaders()
.getContentDisposition()
@@ -759,12 +790,13 @@ class PDFToFileTest {
return mockExecutorResult;
});
ResponseEntity<byte[]> response =
ResponseEntity<StreamingResponseBody> response =
pdfToFileWithUno.processPdfToOfficeFormat(pdfFile, "docx", "writer_pdf_import");
assertEquals(HttpStatus.OK, response.getStatusCode());
assertNotNull(response.getBody());
assertTrue(response.getBody().length > 0);
byte[] bodyBytes = drain(response);
assertNotNull(bodyBytes);
assertTrue(bodyBytes.length > 0);
assertTrue(
response.getHeaders()
.getContentDisposition()