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))