70 lines
1.2 KiB
Go
70 lines
1.2 KiB
Go
package testsupport
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"log"
|
|
)
|
|
|
|
type ChannelReadWriter struct {
|
|
ctx context.Context
|
|
receiver <-chan []byte
|
|
// bytes that were read and that did not fit
|
|
readBuf []byte
|
|
sender chan<- []byte
|
|
closed bool
|
|
}
|
|
|
|
func NewChannelReadWriter(ctx context.Context, receiver <-chan []byte,
|
|
sender chan<- []byte) *ChannelReadWriter {
|
|
return &ChannelReadWriter{
|
|
ctx: ctx,
|
|
receiver: receiver,
|
|
sender: sender,
|
|
closed: false,
|
|
}
|
|
}
|
|
|
|
func (rw *ChannelReadWriter) Read(p []byte) (n int, err error) {
|
|
nread := copy(p, rw.readBuf)
|
|
if nread > 0 {
|
|
log.Printf("Read %v bytes", nread)
|
|
rw.readBuf = rw.readBuf[nread:]
|
|
return nread, nil
|
|
}
|
|
|
|
select {
|
|
case <-rw.ctx.Done():
|
|
return 0, io.EOF
|
|
case data, ok := <-rw.receiver:
|
|
if !ok {
|
|
return 0, io.EOF
|
|
}
|
|
nread = copy(p, data)
|
|
rw.readBuf = data[nread:]
|
|
return nread, nil
|
|
}
|
|
|
|
}
|
|
func (rw *ChannelReadWriter) Write(p []byte) (n int, err error) {
|
|
if rw.closed {
|
|
return 0, errors.New("Write on closed channel")
|
|
}
|
|
select {
|
|
case <-rw.ctx.Done():
|
|
rw.Close()
|
|
return 0, io.EOF
|
|
case rw.sender <- p:
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
func (rw *ChannelReadWriter) Close() error {
|
|
if !rw.closed {
|
|
close(rw.sender)
|
|
rw.closed = true
|
|
}
|
|
return nil
|
|
}
|