diff --git a/node.go b/node.go
index d025f33..5d5a560 100644
--- a/node.go
+++ b/node.go
@@ -377,7 +377,11 @@ func (n *Node) WriteWithOptions(writer io.Writer, opts ...OutputOption) (err err
pastPreserveSpaces := config.preserveSpaces
preserveSpaces := calculatePreserveSpaces(n, pastPreserveSpaces)
b := bufio.NewWriter(writer)
- defer b.Flush()
+ defer func() {
+ if flushErr := b.Flush(); err == nil {
+ err = flushErr
+ }
+ }()
ident := newIndentation(config.useIndentation, b)
if config.printSelf && n.Type != DocumentNode {
diff --git a/node_test.go b/node_test.go
index 070fc8c..7a6c632 100644
--- a/node_test.go
+++ b/node_test.go
@@ -2,7 +2,9 @@ package xmlquery
import (
"encoding/xml"
+ "errors"
"html"
+ "io"
"reflect"
"strings"
"testing"
@@ -564,6 +566,93 @@ func TestWriteWithNamespacePrefix(t *testing.T) {
}
}
+type xmlWriterFunc func([]byte) (int, error)
+
+func (w xmlWriterFunc) Write(p []byte) (int, error) {
+ return w(p)
+}
+
+func TestWriteErrors(t *testing.T) {
+ writeErr := errors.New("write failed")
+ for _, size := range []struct {
+ name string
+ text string
+ }{
+ {"buffered", "value"},
+ {"large", strings.Repeat("x", 8192)},
+ } {
+ t.Run(size.name, func(t *testing.T) {
+ doc, err := Parse(strings.NewReader("" + size.text + ""))
+ if err != nil {
+ t.Fatal(err)
+ }
+ root := doc.SelectElement("root")
+ if root == nil {
+ t.Fatal("root element not found")
+ }
+ children := "" + size.text + ""
+ whole := "" + children + ""
+ declaration := ``
+ indentedChildren := "\n" + size.text + "\n"
+ indentedWhole := "\n\n " + size.text + "\n \n"
+ for _, mode := range []struct {
+ name string
+ write func(io.Writer) error
+ want string
+ }{
+ {"document self", func(w io.Writer) error { return doc.Write(w, true) }, declaration + whole},
+ {"document children", func(w io.Writer) error { return doc.Write(w, false) }, declaration + whole},
+ {"element self", func(w io.Writer) error { return root.Write(w, true) }, whole},
+ {"element children", func(w io.Writer) error { return root.Write(w, false) }, children},
+ {"indented document", func(w io.Writer) error { return doc.WriteWithOptions(w, WithIndentation(" ")) }, declaration + indentedWhole},
+ {"indented element", func(w io.Writer) error { return root.WriteWithOptions(w, WithOutputSelf(), WithIndentation(" ")) }, indentedWhole},
+ {"indented children", func(w io.Writer) error { return root.WriteWithOptions(w, WithIndentation(" ")) }, indentedChildren},
+ } {
+ t.Run(mode.name, func(t *testing.T) {
+ for _, writer := range []struct {
+ name string
+ limit int
+ err error
+ wantErr error
+ }{
+ {"success", -1, nil, nil},
+ {"error", 0, writeErr, writeErr},
+ {"partial error", 2, writeErr, writeErr},
+ {"short write", 2, nil, io.ErrShortWrite},
+ } {
+ t.Run(writer.name, func(t *testing.T) {
+ var output strings.Builder
+ calls := 0
+ err := mode.write(xmlWriterFunc(func(p []byte) (int, error) {
+ calls++
+ n := len(p)
+ if writer.limit >= 0 && n > writer.limit {
+ n = writer.limit
+ }
+ output.Write(p[:n])
+ return n, writer.err
+ }))
+ if err != writer.wantErr {
+ t.Fatalf("expected error %v, got %v", writer.wantErr, err)
+ }
+ want := mode.want
+ if writer.wantErr != nil {
+ want = want[:writer.limit]
+ if calls != 1 {
+ t.Errorf("expected one write attempt, got %d", calls)
+ }
+ }
+ if output.String() != want {
+ t.Errorf("expected output %q, got %q", want, output.String())
+ }
+ })
+ }
+ })
+ }
+ })
+ }
+}
+
func TestQueryWithPrefix(t *testing.T) {
s := `ns2:ClientThis is a client fault`
doc, _ := Parse(strings.NewReader(s))