influxdb/write/batcher_test.go

430 lines
9.7 KiB
Go

package write
import (
"bufio"
"context"
"fmt"
"io"
"io/ioutil"
"strings"
"testing"
"time"
"github.com/google/go-cmp/cmp"
platform "github.com/influxdata/influxdb"
"github.com/influxdata/influxdb/mock"
)
func TestScanLines(t *testing.T) {
tests := []struct {
name string
input string
want []string
wantErr bool
}{
{
name: "3 lines produced including their newlines",
input: "m1,t1=v1 f1=1\nm2,t2=v2 f2=2\nm3,t3=v3 f3=3",
want: []string{"m1,t1=v1 f1=1\n", "m2,t2=v2 f2=2\n", "m3,t3=v3 f3=3"},
},
{
name: "single line without newline",
input: "m1,t1=v1 f1=1",
want: []string{"m1,t1=v1 f1=1"},
},
{
name: "single line with newline",
input: "m1,t1=v1 f1=1\n",
want: []string{"m1,t1=v1 f1=1\n"},
},
{
name: "no lines",
input: "",
want: []string{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
scanner := bufio.NewScanner(strings.NewReader(tt.input))
scanner.Split(ScanLines)
got := []string{}
for scanner.Scan() {
got = append(got, scanner.Text())
}
err := scanner.Err()
if (err != nil) != tt.wantErr {
t.Errorf("ScanLines() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !cmp.Equal(got, tt.want) {
t.Errorf("%q. ScanLines() = -got/+want %s", tt.name, cmp.Diff(got, tt.want))
}
})
}
}
// errorReader mocks io.Reader but returns an error.
type errorReader struct{}
func (r *errorReader) Read(p []byte) (n int, err error) {
return 0, fmt.Errorf("error")
}
func TestBatcher_read(t *testing.T) {
type args struct {
cancel bool
r io.Reader
lines chan []byte
errC chan error
}
tests := []struct {
name string
args args
want []string
wantErr bool
}{
{
name: "reading two lines produces 2 lines",
args: args{
r: strings.NewReader("m1,t1=v1 f1=1\nm2,t2=v2 f2=2"),
lines: make(chan []byte),
errC: make(chan error, 1),
},
want: []string{"m1,t1=v1 f1=1\n", "m2,t2=v2 f2=2"},
},
{
name: "canceling returns no lines",
args: args{
cancel: true,
r: strings.NewReader("m1,t1=v1 f1=1"),
lines: make(chan []byte),
errC: make(chan error, 1),
},
want: []string{},
wantErr: true,
},
{
name: "error from reader returns error",
args: args{
r: &errorReader{},
lines: make(chan []byte),
errC: make(chan error, 1),
},
want: []string{},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
var cancel context.CancelFunc
if tt.args.cancel {
ctx, cancel = context.WithCancel(ctx)
}
b := &Batcher{}
got := []string{}
go b.read(ctx, tt.args.r, tt.args.lines, tt.args.errC)
if cancel != nil {
cancel()
} else {
for line := range tt.args.lines {
got = append(got, string(line))
}
}
err := <-tt.args.errC
if (err != nil) != tt.wantErr {
t.Errorf("ScanLines.read() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !cmp.Equal(got, tt.want) {
t.Errorf("%q. Batcher.read() = -got/+want %s", tt.name, cmp.Diff(got, tt.want))
}
})
}
}
func TestBatcher_write(t *testing.T) {
type fields struct {
MaxFlushBytes int
MaxFlushInterval time.Duration
}
type args struct {
cancel bool
writeError bool
org platform.ID
bucket platform.ID
line string
lines chan []byte
errC chan error
}
tests := []struct {
name string
fields fields
args args
want string
wantErr bool
}{
{
name: "sending a single line will send a line to the service",
fields: fields{
MaxFlushBytes: 1,
},
args: args{
org: platform.ID(1),
bucket: platform.ID(2),
line: "m1,t1=v1 f1=1",
lines: make(chan []byte),
errC: make(chan error),
},
want: "m1,t1=v1 f1=1",
},
{
name: "sending a single line and waiting for a timeout will send a line to the service",
fields: fields{
MaxFlushInterval: time.Millisecond,
},
args: args{
org: platform.ID(1),
bucket: platform.ID(2),
line: "m1,t1=v1 f1=1",
lines: make(chan []byte),
errC: make(chan error),
},
want: "m1,t1=v1 f1=1",
},
{
name: "write service returning error stops the write after timeout",
fields: fields{
MaxFlushInterval: time.Millisecond,
},
args: args{
writeError: true,
org: platform.ID(1),
bucket: platform.ID(2),
line: "m1,t1=v1 f1=1",
lines: make(chan []byte),
errC: make(chan error),
},
wantErr: true,
},
{
name: "canceling will write no data to service",
fields: fields{
MaxFlushBytes: 1,
},
args: args{
cancel: true,
org: platform.ID(1),
bucket: platform.ID(2),
line: "m1,t1=v1 f1=1",
lines: make(chan []byte),
errC: make(chan error),
},
wantErr: true,
},
{
name: "write service returning error stops the write",
fields: fields{
MaxFlushBytes: 1,
},
args: args{
writeError: true,
org: platform.ID(1),
bucket: platform.ID(2),
line: "m1,t1=v1 f1=1",
lines: make(chan []byte),
errC: make(chan error),
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
ctx := context.Background()
var cancel context.CancelFunc
if tt.args.cancel {
ctx, cancel = context.WithCancel(ctx)
}
// mocking the write service here to either return an error
// or get back all the bytes from the reader.
var got string
svc := &mock.WriteService{
WriteF: func(ctx context.Context, org, bucket platform.ID, r io.Reader) error {
if tt.args.writeError {
return fmt.Errorf("error")
}
b, err := ioutil.ReadAll(r)
got = string(b)
return err
},
}
b := &Batcher{
MaxFlushBytes: tt.fields.MaxFlushBytes,
MaxFlushInterval: tt.fields.MaxFlushInterval,
Service: svc,
}
go b.write(ctx, tt.args.org, tt.args.bucket, tt.args.lines, tt.args.errC)
if cancel != nil {
cancel()
} else {
tt.args.lines <- []byte(tt.args.line)
// if the max flush interval is not zero, we are testing to see
// if the data is flushed via the timer rather than forced by
// closing the channel.
if tt.fields.MaxFlushInterval != 0 {
time.Sleep(tt.fields.MaxFlushInterval * 100)
}
close(tt.args.lines)
}
err := <-tt.args.errC
if (err != nil) != tt.wantErr {
t.Errorf("ScanLines.read() error = %v, wantErr %v", err, tt.wantErr)
return
}
if !cmp.Equal(got, tt.want) {
t.Errorf("%q. Batcher.write() = -got/+want %s", tt.name, cmp.Diff(got, tt.want))
}
})
}
}
func TestBatcher_Write(t *testing.T) {
type fields struct {
MaxFlushBytes int
MaxFlushInterval time.Duration
}
type args struct {
writeError bool
org platform.ID
bucket platform.ID
r io.Reader
}
tests := []struct {
name string
fields fields
args args
want string
wantFlushes int
wantErr bool
}{
{
name: "a line of line protocol is sent to the service",
fields: fields{
MaxFlushBytes: 1,
},
args: args{
org: platform.ID(1),
bucket: platform.ID(2),
r: strings.NewReader("m1,t1=v1 f1=1"),
},
want: "m1,t1=v1 f1=1",
wantFlushes: 1,
},
{
name: "multiple lines cause multiple flushes",
fields: fields{
MaxFlushBytes: len([]byte("m1,t1=v1 f1=1\n")),
},
args: args{
org: platform.ID(1),
bucket: platform.ID(2),
r: strings.NewReader("m1,t1=v1 f1=1\nm2,t2=v2 f2=2\nm3,t3=v3 f3=3"),
},
want: "m3,t3=v3 f3=3",
wantFlushes: 3,
},
{
name: "errors during read return error",
fields: fields{},
args: args{
org: platform.ID(1),
bucket: platform.ID(2),
r: &errorReader{},
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// mocking the write service here to either return an error
// or get back all the bytes from the reader.
var (
got string
gotFlushes int
)
svc := &mock.WriteService{
WriteF: func(ctx context.Context, org, bucket platform.ID, r io.Reader) error {
if tt.args.writeError {
return fmt.Errorf("error")
}
b, err := ioutil.ReadAll(r)
got = string(b)
gotFlushes++
return err
},
}
b := &Batcher{
MaxFlushBytes: tt.fields.MaxFlushBytes,
MaxFlushInterval: tt.fields.MaxFlushInterval,
Service: svc,
}
ctx := context.Background()
if err := b.Write(ctx, tt.args.org, tt.args.bucket, tt.args.r); (err != nil) != tt.wantErr {
t.Errorf("Batcher.Write() error = %v, wantErr %v", err, tt.wantErr)
}
if gotFlushes != tt.wantFlushes {
t.Errorf("%q. Batcher.Write() flushes %d want %d", tt.name, gotFlushes, tt.wantFlushes)
}
if !cmp.Equal(got, tt.want) {
t.Errorf("%q. Batcher.Write() = -got/+want %s", tt.name, cmp.Diff(got, tt.want))
}
})
}
}
func TestBatcher_WriteTimeout(t *testing.T) {
// mocking the write service here to either return an error
// or get back all the bytes from the reader.
var got string
svc := &mock.WriteService{
WriteF: func(ctx context.Context, org, bucket platform.ID, r io.Reader) error {
b, err := ioutil.ReadAll(r)
got = string(b)
return err
},
}
b := &Batcher{
Service: svc,
}
// this mimics a reader like stdin that may never return data.
r, _ := io.Pipe()
ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
defer cancel()
if err := b.Write(ctx, platform.ID(1), platform.ID(2), r); err != context.DeadlineExceeded {
t.Errorf("Batcher.Write() with timeout error = %v", err)
}
if got != "" {
t.Errorf(" Batcher.Write() with timeout got %s", got)
}
}