| 12 | ) |
| 13 | |
| 14 | func TestPDFLoader(t *testing.T) { |
| 15 | ctx := context.Background() |
| 16 | t.Parallel() |
| 17 | |
| 18 | page1Content := " A Simple PDF File This is a small demonstration .pdf file - " + |
| 19 | "just for use in the Virtual Mechanics tutorials. More text. And more text. And more " + |
| 20 | "text. And more text. And more text. And more text. And more text. And more text. " + |
| 21 | "And more text. And more text. And more text. Boring, zzzzz. And more text. And more " + |
| 22 | "text. And more text. And more text. And more text. And more text. And more text. " + |
| 23 | "And more text. And more text. And more text. And more text. And more text. And more " + |
| 24 | "text. And more text. And more text. And more text. Even more. Continued on page 2 ..." |
| 25 | |
| 26 | page2Content := " Simple PDF File 2 ...continued from page 1. Yet more text. And more " + |
| 27 | "text. And more text. And more text. And more text. And more text. And more text. And more " + |
| 28 | " text. Oh, how boring typing this stuff. But not as boring as watching paint dry. And more " + |
| 29 | "text. And more text. And more text. And more text. Boring. More, a little more text. " + |
| 30 | "The end, and just as well. " |
| 31 | |
| 32 | expectedResults := []struct { |
| 33 | content string |
| 34 | metadata map[string]any |
| 35 | }{ |
| 36 | {content: page1Content, metadata: map[string]any{"page": 1, "total_pages": 2}}, |
| 37 | {content: page2Content, metadata: map[string]any{"page": 2, "total_pages": 2}}, |
| 38 | } |
| 39 | |
| 40 | t.Run("PDFLoad", func(t *testing.T) { |
| 41 | t.Parallel() |
| 42 | f, err := os.Open("./testdata/sample.pdf") |
| 43 | require.NoError(t, err) |
| 44 | defer f.Close() |
| 45 | finfo, err := f.Stat() |
| 46 | require.NoError(t, err) |
| 47 | p := NewPDF(f, finfo.Size()) |
| 48 | docs, err := p.Load(ctx) |
| 49 | require.NoError(t, err) |
| 50 | |
| 51 | assert.Len(t, docs, 2) |
| 52 | |
| 53 | for r := range expectedResults { |
| 54 | assert.Equal(t, expectedResults[r].content, docs[r].PageContent) |
| 55 | assert.Equal(t, expectedResults[r].metadata, docs[r].Metadata) |
| 56 | } |
| 57 | }) |
| 58 | |
| 59 | t.Run("PDFLoadPassword", func(t *testing.T) { |
| 60 | t.Parallel() |
| 61 | f, err := os.Open("./testdata/sample_password.pdf") |
| 62 | require.NoError(t, err) |
| 63 | defer f.Close() |
| 64 | finfo, err := f.Stat() |
| 65 | require.NoError(t, err) |
| 66 | p := NewPDF(f, finfo.Size(), WithPassword("password")) |
| 67 | docs, err := p.Load(ctx) |
| 68 | require.NoError(t, err) |
| 69 | |
| 70 | assert.Len(t, docs, 2) |
| 71 | |