Skip to content

Commit 089c3eb

Browse files
JM-Labilayaperumalg
authored andcommitted
Add FileDocumentWriterTest for FileDocumentWriter default constructor and various options
Auto-cherry-pick to 1.0.x Fixes #3646 Signed-off-by: Jemin Huh <[email protected]>
1 parent 28f86bf commit 089c3eb

File tree

1 file changed

+145
-0
lines changed

1 file changed

+145
-0
lines changed
Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
package org.springframework.ai.writer;
2+
3+
import org.junit.jupiter.api.BeforeEach;
4+
import org.junit.jupiter.api.Test;
5+
import org.junit.jupiter.api.io.TempDir;
6+
import org.springframework.ai.document.Document;
7+
import org.springframework.ai.document.MetadataMode;
8+
9+
import java.io.IOException;
10+
import java.nio.file.Files;
11+
import java.nio.file.Path;
12+
import java.util.List;
13+
14+
import static org.junit.jupiter.api.Assertions.*;
15+
16+
/**
17+
* @author Jemin Huh
18+
*/
19+
public class FileDocumentWriterTest {
20+
21+
@TempDir
22+
Path tempDir;
23+
24+
private String testFileName;
25+
26+
private List<Document> testDocuments;
27+
28+
@BeforeEach
29+
void setUp() {
30+
testFileName = tempDir.resolve("file-document-test-output.txt").toString();
31+
testDocuments = List.of(
32+
Document.builder()
33+
.text("Document one introduces the core functionality of Spring AI.")
34+
.metadata("page_number", "1")
35+
.metadata("end_page_number", "2")
36+
.metadata("source", "intro.pdf")
37+
.metadata("title", "Spring AI Overview")
38+
.metadata("author", "QA Team")
39+
.build(),
40+
Document.builder()
41+
.text("Document two illustrates multi-line handling and line breaks.\nEnsure preservation of formatting.")
42+
.metadata("page_number", "3")
43+
.metadata("end_page_number", "4")
44+
.metadata("source", "formatting.pdf")
45+
.build(),
46+
Document.builder()
47+
.text("Document three checks metadata inclusion and output formatting behavior.")
48+
.metadata("page_number", "5")
49+
.metadata("end_page_number", "6")
50+
.metadata("version", "v1.2")
51+
.build());
52+
}
53+
54+
@Test
55+
void testBasicWrite() throws IOException {
56+
var writer = new FileDocumentWriter(testFileName);
57+
writer.accept(testDocuments);
58+
59+
List<String> lines = Files.readAllLines(Path.of(testFileName));
60+
assertEquals("", lines.get(0));
61+
assertEquals("", lines.get(1));
62+
assertEquals("Document one introduces the core functionality of Spring AI.", lines.get(2));
63+
assertEquals("", lines.get(3));
64+
assertEquals("Document two illustrates multi-line handling and line breaks.", lines.get(4));
65+
assertEquals("Ensure preservation of formatting.", lines.get(5));
66+
assertEquals("", lines.get(6));
67+
assertEquals("Document three checks metadata inclusion and output formatting behavior.", lines.get(7));
68+
}
69+
70+
@Test
71+
void testWriteWithDocumentMarkers() throws IOException {
72+
var writer = new FileDocumentWriter(testFileName, true, MetadataMode.NONE, false);
73+
writer.accept(testDocuments);
74+
75+
List<String> lines = Files.readAllLines(Path.of(testFileName));
76+
assertEquals("", lines.get(0));
77+
assertEquals("### Doc: 0, pages:[1,2]", lines.get(1));
78+
assertEquals("", lines.get(2));
79+
assertEquals("", lines.get(3));
80+
assertEquals("Document one introduces the core functionality of Spring AI.", lines.get(4));
81+
assertEquals("### Doc: 1, pages:[3,4]", lines.get(5));
82+
assertEquals("", lines.get(6));
83+
assertEquals("", lines.get(7));
84+
assertEquals("Document two illustrates multi-line handling and line breaks.", lines.get(8));
85+
assertEquals("Ensure preservation of formatting.", lines.get(9));
86+
assertEquals("### Doc: 2, pages:[5,6]", lines.get(10));
87+
assertEquals("", lines.get(11));
88+
assertEquals("", lines.get(12));
89+
assertEquals("Document three checks metadata inclusion and output formatting behavior.", lines.get(13));
90+
}
91+
92+
@Test
93+
void testMetadataModeAllWithDocumentMarkers() throws IOException {
94+
var writer = new FileDocumentWriter(testFileName, true, MetadataMode.ALL, false);
95+
writer.accept(testDocuments);
96+
97+
List<String> lines = Files.readAllLines(Path.of(testFileName));
98+
assertEquals("", lines.get(0));
99+
assertEquals("### Doc: 0, pages:[1,2]", lines.get(1));
100+
String subListToString = lines.subList(2, 7).toString();
101+
assertTrue(subListToString.contains("page_number: 1"));
102+
assertTrue(subListToString.contains("end_page_number: 2"));
103+
assertTrue(subListToString.contains("source: intro.pdf"));
104+
assertTrue(subListToString.contains("title: Spring AI Overview"));
105+
assertTrue(subListToString.contains("author: QA Team"));
106+
assertEquals("", lines.get(7));
107+
assertEquals("Document one introduces the core functionality of Spring AI.", lines.get(8));
108+
109+
assertEquals("### Doc: 1, pages:[3,4]", lines.get(9));
110+
subListToString = lines.subList(10, 13).toString();
111+
assertTrue(subListToString.contains("page_number: 3"));
112+
assertTrue(subListToString.contains("source: formatting.pdf"));
113+
assertTrue(subListToString.contains("end_page_number: 4"));
114+
assertEquals("", lines.get(13));
115+
assertEquals("Document two illustrates multi-line handling and line breaks.", lines.get(14));
116+
assertEquals("Ensure preservation of formatting.", lines.get(15));
117+
118+
assertEquals("### Doc: 2, pages:[5,6]", lines.get(16));
119+
subListToString = lines.subList(17, 20).toString();
120+
assertTrue(subListToString.contains("page_number: 5"));
121+
assertTrue(subListToString.contains("end_page_number: 6"));
122+
assertTrue(subListToString.contains("version: v1.2"));
123+
assertEquals("", lines.get(20));
124+
assertEquals("Document three checks metadata inclusion and output formatting behavior.", lines.get(21));
125+
}
126+
127+
@Test
128+
void testAppendWrite() throws IOException {
129+
Files.writeString(Path.of(testFileName), "Test String\n");
130+
131+
var writer = new FileDocumentWriter(testFileName, false, MetadataMode.NONE, true);
132+
writer.accept(testDocuments.subList(0, 2));
133+
134+
List<String> lines = Files.readAllLines(Path.of(testFileName));
135+
assertEquals("Test String", lines.get(0));
136+
assertEquals("", lines.get(1));
137+
assertEquals("", lines.get(2));
138+
assertEquals("Document one introduces the core functionality of Spring AI.", lines.get(3));
139+
assertEquals("", lines.get(4));
140+
assertEquals("Document two illustrates multi-line handling and line breaks.", lines.get(5));
141+
assertEquals("Ensure preservation of formatting.", lines.get(6));
142+
assertEquals(7, lines.size());
143+
}
144+
145+
}

0 commit comments

Comments
 (0)