diff --git a/example-exports/examplegen.go b/example-exports/examplegen.go index b156167..e1e3deb 100644 --- a/example-exports/examplegen.go +++ b/example-exports/examplegen.go @@ -80,8 +80,11 @@ func main() { log.Panic(errors.Wrapf(err, "write message %q", msg)) } } - if _, err := of.Flush(); err != nil { + if _, err := of.Stage(); err != nil { log.Panic(errors.Wrap(err, "stage outfile")) } + if err := of.Flush(); err != nil { + log.Panic(errors.Wrap(err, "flush outfile")) + } } } diff --git a/internal/bagoup/export_test.go b/internal/bagoup/export_test.go index 087e9dc..196011a 100644 --- a/internal/bagoup/export_test.go +++ b/internal/bagoup/export_test.go @@ -65,14 +65,16 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid;;;testguid2.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), dbMock.EXPECT().GetMessageIDs(3), osMock.EXPECT().MkdirAll("messages-export/testdisplayname2", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname2/testguid3.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[1]), - ofMocks[1].EXPECT().Flush(), + ofMocks[1].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[1].EXPECT().Flush(), ) }, }, @@ -113,8 +115,9 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid;;;testguid2.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), ) }, }, @@ -155,14 +158,16 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid;;;testguid2.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), dbMock.EXPECT().GetMessageIDs(3), osMock.EXPECT().MkdirAll("messages-export/testdisplayname2", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname2/testguid3.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[1]), - ofMocks[1].EXPECT().Flush(), + ofMocks[1].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[1].EXPECT().Flush(), ) }, }, @@ -193,14 +198,16 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), dbMock.EXPECT().GetMessageIDs(2), osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid2.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[1]), - ofMocks[1].EXPECT().Flush(), + ofMocks[1].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[1].EXPECT().Flush(), ) }, }, @@ -228,8 +235,9 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid.pdf").Return(chatFile, nil), osMock.EXPECT().NewPDFOutFile(chatFile, gomock.Any(), false).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), ) }, }, @@ -257,8 +265,9 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid.pdf").Return(chatFile, nil), osMock.EXPECT().NewPDFOutFile(chatFile, gomock.Any(), false).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), ) }, }, @@ -287,8 +296,9 @@ func TestExportChats(t *testing.T) { osMock.EXPECT().MkdirAll("messages-export/testdisplayname/attachments", os.ModePerm), osMock.EXPECT().Create("messages-export/testdisplayname/testguid.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMocks[0]), - ofMocks[0].EXPECT().Flush(), + ofMocks[0].EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMocks[0].EXPECT().Flush(), ) }, }, diff --git a/internal/bagoup/write.go b/internal/bagoup/write.go index 77e517b..14da7cc 100644 --- a/internal/bagoup/write.go +++ b/internal/bagoup/write.go @@ -119,15 +119,18 @@ func (cfg *configuration) handleFileContents(outFile opsys.OutFile, messageIDs [ invalidCount++ } } - imgCount, err := outFile.Flush() + imgCount, err := outFile.Stage() if err != nil { - return errors.Wrapf(err, "flush chat file %q to disk", outFile.Name()) + return errors.Wrapf(err, "stage chat file %q for writing", outFile.Name()) } if openFilesLimit := cfg.OS.GetOpenFilesLimit(); imgCount*2 > openFilesLimit { if err := cfg.OS.SetOpenFilesLimit(imgCount * 2); err != nil { return errors.Wrapf(err, "chat file %q - increase the open file limit from %d to %d to support %d embedded images", outFile.Name(), openFilesLimit, imgCount*2, imgCount) } } + if err := outFile.Flush(); err != nil { + return errors.Wrapf(err, "flush chat file %q to disk", outFile.Name()) + } cfg.counts.files++ cfg.counts.messages += msgCount cfg.counts.messagesInvalid += invalidCount diff --git a/internal/bagoup/write_test.go b/internal/bagoup/write_test.go index 3f94e7b..6b117cd 100644 --- a/internal/bagoup/write_test.go +++ b/internal/bagoup/write_test.go @@ -53,8 +53,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.txt"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 1, @@ -79,8 +80,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 2, @@ -107,9 +109,10 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush().Return(500, nil), + ofMock.EXPECT().Stage().Return(500, nil), osMock.EXPECT().GetOpenFilesLimit().Return(256), osMock.EXPECT().SetOpenFilesLimit(1000), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 2, @@ -136,8 +139,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.txt"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 1, @@ -165,8 +169,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2-1.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 1, @@ -244,7 +249,7 @@ func TestWriteFile(t *testing.T) { wantErr: `write message "message2" to file "messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.txt": this is an outfile error`, }, { - msg: "Flush error", + msg: "Staging error", pdf: true, setupMocks: func(dbMock *mock_chatdb.MockChatDB, osMock *mock_opsys.MockOS, ofMock *mock_opsys.MockOutFile) { gomock.InOrder( @@ -263,11 +268,11 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush().Return(0, errors.New("this is a staging error")), + ofMock.EXPECT().Stage().Return(0, errors.New("this is a staging error")), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ) }, - wantErr: `flush chat file "messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf" to disk: this is a staging error`, + wantErr: `stage chat file "messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf" for writing: this is a staging error`, }, { msg: "open files limit increase fails", @@ -289,7 +294,7 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush().Return(500, nil), + ofMock.EXPECT().Stage().Return(500, nil), osMock.EXPECT().GetOpenFilesLimit().Return(256), osMock.EXPECT().SetOpenFilesLimit(1000).Return(errors.New("this is a syscall error")), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), @@ -297,6 +302,34 @@ func TestWriteFile(t *testing.T) { }, wantErr: `chat file "messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf" - increase the open file limit from 256 to 1000 to support 500 embedded images: this is a syscall error`, }, + { + msg: "Flush error", + pdf: true, + setupMocks: func(dbMock *mock_chatdb.MockChatDB, osMock *mock_opsys.MockOS, ofMock *mock_opsys.MockOutFile) { + gomock.InOrder( + osMock.EXPECT().MkdirAll("messages-export/friend", os.ModePerm), + osMock.EXPECT().Create("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf").Return(chatFile, nil), + osMock.EXPECT().NewPDFOutFile(chatFile, gomock.Any(), false).Return(ofMock), + dbMock.EXPECT().GetMessage(1, nil).Return("message1", true, nil), + ofMock.EXPECT().WriteMessage("message1"), + dbMock.EXPECT().GetMessage(2, nil).Return("message2", true, nil), + ofMock.EXPECT().WriteMessage("message2"), + osMock.EXPECT().FileExist("attachment1.heic").Return(true, nil), + osMock.EXPECT().HEIC2JPG("attachment1.heic").Return("attachment1.jpeg", nil), + ofMock.EXPECT().WriteAttachment("attachment1.jpeg"), + osMock.EXPECT().FileExist("attachment2.jpeg").Return(true, nil), + osMock.EXPECT().HEIC2JPG("attachment2.jpeg").Return("attachment2.jpeg", nil), + ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), + ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), + ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), + ofMock.EXPECT().Stage(), + osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush().Return(errors.New("this is a flush error")), + ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), + ) + }, + wantErr: `flush chat file "messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf" to disk: this is a flush error`, + }, { msg: "attachment file does not exist", setupMocks: func(dbMock *mock_chatdb.MockChatDB, osMock *mock_opsys.MockOS, ofMock *mock_opsys.MockOutFile) { @@ -315,8 +348,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.txt"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 1, @@ -421,8 +455,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.pdf"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantJPGs: 1, @@ -465,8 +500,9 @@ func TestWriteFile(t *testing.T) { ofMock.EXPECT().WriteAttachment("attachment2.jpeg"), ofMock.EXPECT().Name().Return("messages-export/friend/iMessage;-;friend@gmail.com;;;iMessage;-;friend@hotmail.com.txt"), ofMock.EXPECT().ReferenceAttachment("att3transfer.png"), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) }, wantInvalid: 1, @@ -538,8 +574,9 @@ func TestWriteFile(t *testing.T) { osMock.EXPECT().MkdirAll("friend", os.ModePerm), osMock.EXPECT().Create("friend/iMessage;-;heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress.heresareallylongemailaddress@gmail.c.txt").Return(chatFile, nil), osMock.EXPECT().NewTxtOutFile(chatFile).Return(ofMock), - ofMock.EXPECT().Flush(), + ofMock.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock.EXPECT().Flush(), ) cfg := configuration{OS: osMock} @@ -579,8 +616,9 @@ func TestWriteFile(t *testing.T) { } mockCalls = append( mockCalls, - ofMock1.EXPECT().Flush(), + ofMock1.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock1.EXPECT().Flush(), osMock.EXPECT().Create("messages-export/friend/iMessage;-;friend@gmail.com.2.pdf").Return(chatFile2, nil), osMock.EXPECT().NewPDFOutFile(chatFile2, gomock.Any(), false).Return(ofMock2), ) @@ -595,8 +633,9 @@ func TestWriteFile(t *testing.T) { } mockCalls = append( mockCalls, - ofMock2.EXPECT().Flush(), + ofMock2.EXPECT().Stage(), osMock.EXPECT().GetOpenFilesLimit().Return(256), + ofMock2.EXPECT().Flush(), ) gomock.InOrder(mockCalls...) diff --git a/opsys/mock_opsys/mock_outfile.go b/opsys/mock_opsys/mock_outfile.go index 2ff53bb..f4ce6b5 100644 --- a/opsys/mock_opsys/mock_outfile.go +++ b/opsys/mock_opsys/mock_outfile.go @@ -34,12 +34,11 @@ func (m *MockOutFile) EXPECT() *MockOutFileMockRecorder { } // Flush mocks base method. -func (m *MockOutFile) Flush() (int, error) { +func (m *MockOutFile) Flush() error { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "Flush") - ret0, _ := ret[0].(int) - ret1, _ := ret[1].(error) - return ret0, ret1 + ret0, _ := ret[0].(error) + return ret0 } // Flush indicates an expected call of Flush. @@ -76,6 +75,21 @@ func (mr *MockOutFileMockRecorder) ReferenceAttachment(arg0 interface{}) *gomock return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReferenceAttachment", reflect.TypeOf((*MockOutFile)(nil).ReferenceAttachment), arg0) } +// Stage mocks base method. +func (m *MockOutFile) Stage() (int, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Stage") + ret0, _ := ret[0].(int) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Stage indicates an expected call of Stage. +func (mr *MockOutFileMockRecorder) Stage() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stage", reflect.TypeOf((*MockOutFile)(nil).Stage)) +} + // WriteAttachment mocks base method. func (m *MockOutFile) WriteAttachment(arg0 string) (bool, error) { m.ctrl.T.Helper() diff --git a/opsys/outfile.go b/opsys/outfile.go index 112985d..c540dad 100644 --- a/opsys/outfile.go +++ b/opsys/outfile.go @@ -55,9 +55,11 @@ type ( WriteAttachment(attPath string) (bool, error) // ReferenceAttachment adds a reference to the given filename in the Outfile. ReferenceAttachment(filename string) error - // Flush flushes the contents of an OutFile to disk, and returns the - // number of images embedded in the OutFile. - Flush() (int, error) + // Stage prepares the OutFile for flushing to disk, and returns the number + // of images embedded in the OutFile. + Stage() (int, error) + // Flush flushes the contents of an OutFile to disk. + Flush() error } ) @@ -82,10 +84,14 @@ func (f txtFile) ReferenceAttachment(filename string) error { return f.WriteMessage(fmt.Sprintf("\n", filename)) } -func (f txtFile) Flush() (int, error) { +func (f txtFile) Stage() (int, error) { return 0, nil } +func (f txtFile) Flush() error { + return nil +} + type ( pdfFile struct { afero.File @@ -93,7 +99,7 @@ type ( contents htmlFileData embeddableImageTypes []string templatePath string - html template.HTML + buf bytes.Buffer } htmlFileData struct { @@ -158,21 +164,21 @@ func (f *pdfFile) ReferenceAttachment(filename string) error { return nil } -func (f *pdfFile) Flush() (int, error) { +func (f *pdfFile) Stage() (int, error) { tmpl, err := template.ParseFS(_embedFS, f.templatePath) if err != nil { return 0, errors.Wrap(err, "parse HTML template") } - var buf bytes.Buffer - if err := tmpl.Execute(&buf, f.contents); err != nil { + if err := tmpl.Execute(&f.buf, f.contents); err != nil { return 0, errors.Wrap(err, "execute HTML template") } - f.html = template.HTML(buf.String()) - page := wkhtmltopdf.NewPageReader(bytes.NewReader(buf.Bytes())) + htmlStr := template.HTML(f.buf.String()) + return strings.Count(string(htmlStr), ": can't evaluate field InvalidReference in type opsys.htmlFileData`, + wantStageErr: `execute HTML template: template: outfile_html_invalid.tmpl:1:2: executing "outfile_html_invalid.tmpl" at <.InvalidReference>: can't evaluate field InvalidReference in type opsys.htmlFileData`, }, { msg: "PDF creation error", @@ -190,7 +193,50 @@ func TestPDFFile(t *testing.T) { pMock.EXPECT().Create().Return(errors.New("this is a PDF creation error")), ) }, - wantErr: "write out PDF: this is a PDF creation error", + wantHTML: template.HTML( + ` + + + + + testfile + + + + + + + + + + + + test message
+ tennisballs.jpeg
+ <attached: video.mov>
+ <attached: signallogo.pluginPayloadAttachment>
+ + + +`, + ), + wantImgCount: 1, + wantFlushErr: "this is a PDF creation error", }, } @@ -230,14 +276,22 @@ func TestPDFFile(t *testing.T) { assert.Equal(t, tt.includePPA, embedded) // Stage the PDF - imgCount, err := of.Flush() - if tt.wantErr != "" { - assert.Error(t, err, tt.wantErr) + imgCount, err := of.Stage() + if tt.wantStageErr != "" { + assert.Error(t, err, tt.wantStageErr) return } assert.NilError(t, err) assert.Equal(t, tt.wantImgCount, imgCount) - assert.Equal(t, pdf.html, tt.wantHTML) + assert.Equal(t, template.HTML(pdf.buf.String()), tt.wantHTML) + + // Flush the PDF + err = of.Flush() + if tt.wantFlushErr != "" { + assert.Error(t, err, tt.wantFlushErr) + return + } + assert.NilError(t, err) }) } }