go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/logdog/client/butler/stream_test.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package butler
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"io"
    21  	"testing"
    22  	"time"
    23  
    24  	. "github.com/smartystreets/goconvey/convey"
    25  	"go.chromium.org/luci/common/clock"
    26  	"go.chromium.org/luci/common/clock/testclock"
    27  	"go.chromium.org/luci/common/logging"
    28  	"go.chromium.org/luci/logdog/client/butler/bundler"
    29  )
    30  
    31  type testBundlerStream struct {
    32  	closed   bool
    33  	appended []byte
    34  	ts       []time.Time
    35  	err      error
    36  
    37  	data []*testBundlerData
    38  }
    39  
    40  func (bs *testBundlerStream) LeaseData() bundler.Data {
    41  	d := &testBundlerData{
    42  		value: make([]byte, 128),
    43  	}
    44  	bs.data = append(bs.data, d)
    45  	return d
    46  }
    47  
    48  func (bs *testBundlerStream) Append(d bundler.Data) error {
    49  	if bs.err != nil {
    50  		return bs.err
    51  	}
    52  
    53  	tbd := d.(*testBundlerData)
    54  	bs.appended = append(bs.appended, tbd.value...)
    55  	bs.ts = append(bs.ts, d.Timestamp())
    56  	tbd.Release()
    57  	return nil
    58  }
    59  
    60  func (bs *testBundlerStream) Close() {
    61  	if bs.closed {
    62  		panic("double close")
    63  	}
    64  	bs.closed = true
    65  }
    66  
    67  func (bs *testBundlerStream) allReleased() bool {
    68  	for _, d := range bs.data {
    69  		if !d.released {
    70  			return false
    71  		}
    72  	}
    73  	return true
    74  }
    75  
    76  func (bs *testBundlerStream) closedAndReleased() bool {
    77  	return (bs.allReleased() && bs.closed)
    78  }
    79  
    80  type testBundlerData struct {
    81  	value    []byte
    82  	ts       time.Time
    83  	bound    bool
    84  	released bool
    85  }
    86  
    87  func (d *testBundlerData) Release() {
    88  	if d.released {
    89  		panic("double release")
    90  	}
    91  	d.released = true
    92  }
    93  
    94  func (d *testBundlerData) Bytes() []byte {
    95  	return d.value
    96  }
    97  
    98  func (d *testBundlerData) Bind(c int, ts time.Time) bundler.Data {
    99  	d.ts = ts
   100  	d.value = d.value[:c]
   101  	d.bound = true
   102  	return d
   103  }
   104  
   105  func (d *testBundlerData) Timestamp() time.Time {
   106  	if !d.bound {
   107  		panic("not bound, no timestamp")
   108  	}
   109  	return d.ts
   110  }
   111  
   112  type testReadCloser struct {
   113  	data   []byte
   114  	err    error
   115  	closed bool
   116  }
   117  
   118  func (rc *testReadCloser) Read(b []byte) (int, error) {
   119  	if len(b) < len(rc.data) {
   120  		panic("test data too large")
   121  	}
   122  	if rc.closed {
   123  		return 0, io.EOF
   124  	}
   125  	return copy(b, rc.data), rc.err
   126  }
   127  
   128  func (rc *testReadCloser) Close() error {
   129  	if rc.closed {
   130  		panic("double close")
   131  	}
   132  	rc.closed = true
   133  	return rc.err
   134  }
   135  
   136  func TestStream(t *testing.T) {
   137  	Convey(`A testing stream`, t, func() {
   138  		c, tc := testclock.UseTime(context.Background(), testclock.TestTimeUTC)
   139  		bs := &testBundlerStream{}
   140  		rc := &testReadCloser{}
   141  
   142  		s := &stream{
   143  			log: logging.Get(c),
   144  			now: clock.Get(c).Now,
   145  			r:   rc,
   146  			c:   rc,
   147  			bs:  bs,
   148  		}
   149  
   150  		Convey(`Will read chunks until EOF.`, func() {
   151  			rc.data = []byte("foo")
   152  			So(s.readChunk(), ShouldBeTrue)
   153  
   154  			rc.data = []byte(nil)
   155  			So(s.readChunk(), ShouldBeTrue)
   156  
   157  			tc.Add(time.Second)
   158  			rc.data = []byte("bar")
   159  			So(s.readChunk(), ShouldBeTrue)
   160  
   161  			s.closeStream()
   162  			So(s.readChunk(), ShouldBeFalse)
   163  			So(bs.appended, ShouldResemble, []byte("foobar"))
   164  			So(bs.ts, ShouldResemble, []time.Time{testclock.TestTimeUTC, testclock.TestTimeUTC.Add(time.Second)})
   165  			So(bs.closedAndReleased(), ShouldBeTrue)
   166  		})
   167  
   168  		Convey(`Will NOT release Data on Append error.`, func() {
   169  			bs.err = errors.New("test error")
   170  			rc.data = []byte("bar")
   171  			So(s.readChunk(), ShouldBeFalse)
   172  
   173  			s.closeStream()
   174  			So(bs.closed, ShouldBeTrue)
   175  			So(bs.allReleased(), ShouldBeFalse)
   176  		})
   177  
   178  		Convey(`Will halt if a stream error is encountered.`, func() {
   179  			rc.data = []byte("foo")
   180  			rc.err = errors.New("test error")
   181  			So(s.readChunk(), ShouldBeFalse)
   182  
   183  			s.closeStream()
   184  			So(bs.appended, ShouldResemble, []byte("foo"))
   185  			So(bs.closedAndReleased(), ShouldBeTrue)
   186  		})
   187  
   188  		Convey(`Will close Bundler Stream even if Closer returns an error.`, func() {
   189  			rc.err = errors.New("test error")
   190  			s.closeStream()
   191  			So(bs.closedAndReleased(), ShouldBeTrue)
   192  		})
   193  	})
   194  }