Update test.
[arvados.git] / lib / crunchrun / bufthenwrite.go
1 // Copyright (C) The Arvados Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package crunchrun
6
7 import (
8         "bytes"
9         "io"
10         "sync"
11 )
12
13 type bufThenWrite struct {
14         buf bytes.Buffer
15         w   io.Writer
16         mtx sync.Mutex
17 }
18
19 func (btw *bufThenWrite) SetWriter(w io.Writer) error {
20         btw.mtx.Lock()
21         defer btw.mtx.Unlock()
22         btw.w = w
23         _, err := io.Copy(w, &btw.buf)
24         return err
25 }
26
27 func (btw *bufThenWrite) Write(p []byte) (int, error) {
28         btw.mtx.Lock()
29         defer btw.mtx.Unlock()
30         if btw.w == nil {
31                 btw.w = &btw.buf
32         }
33         return btw.w.Write(p)
34 }