package internal_test import ( "bytes" "context" "io" "io/ioutil" "strings" "testing" "github.com/golang/mock/gomock" "github.com/influxdata/influx-cli/v2/internal" "github.com/influxdata/influx-cli/v2/internal/api" "github.com/influxdata/influx-cli/v2/internal/config" "github.com/influxdata/influx-cli/v2/internal/mock" "github.com/stretchr/testify/assert" tmock "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" ) type bufferReader struct { buf bytes.Buffer } func (pr *bufferReader) Open(context.Context) (io.Reader, io.Closer, error) { return &pr.buf, ioutil.NopCloser(nil), nil } type noopThrottler struct { used bool } func (nt *noopThrottler) Throttle(_ context.Context, in io.Reader) io.Reader { nt.used = true return in } type lineBatcher struct{} func (pb *lineBatcher) WriteBatches(_ context.Context, r io.Reader, writeFn func(batch []byte) error) error { buf := bytes.Buffer{} if _, err := io.Copy(&buf, r); err != nil { return err } for _, l := range strings.Split(buf.String(), "\n") { if l != "" { if err := writeFn([]byte(l)); err != nil { return err } } } return nil } func TestWriteByIDs(t *testing.T) { t.Parallel() inLines := []string{"fake line protocol 1", "fake line protocol 2", "fake line protocol 3"} mockReader := bufferReader{} for _, l := range inLines { _, err := io.Copy(&mockReader.buf, strings.NewReader(l+"\n")) require.NoError(t, err) } mockThrottler := noopThrottler{} mockBatcher := lineBatcher{} params := internal.WriteParams{ OrgID: "12345", BucketID: "98765", Precision: api.WRITEPRECISION_S, } cli := internal.CLI{ActiveConfig: config.Config{Org: "my-default-org"}} ctrl := gomock.NewController(t) client := mock.NewMockWriteApi(ctrl) var writtenLines []string client.EXPECT().PostWrite(gomock.Any()).Return(api.ApiPostWriteRequest{ApiService: client}).Times(len(inLines)) client.EXPECT().PostWriteExecute(tmock.MatchedBy(func(in api.ApiPostWriteRequest) bool { return assert.Equal(t, params.OrgID, *in.GetOrg()) && assert.Equal(t, params.BucketID, *in.GetBucket()) && assert.Equal(t, params.Precision, *in.GetPrecision()) && assert.Equal(t, "gzip", *in.GetContentEncoding()) // Make sure the body is properly marked for compression. })).DoAndReturn(func(in api.ApiPostWriteRequest) error { writtenLines = append(writtenLines, string(in.GetBody())) return nil }).Times(len(inLines)) clients := internal.WriteClients{ Reader: &mockReader, Throttler: &mockThrottler, Writer: &mockBatcher, Client: client, } require.NoError(t, cli.Write(context.Background(), &clients, ¶ms)) require.Equal(t, inLines, writtenLines) require.True(t, mockThrottler.used) } func TestWriteByNames(t *testing.T) { t.Parallel() inLines := []string{"fake line protocol 1", "fake line protocol 2", "fake line protocol 3"} mockReader := bufferReader{} for _, l := range inLines { _, err := io.Copy(&mockReader.buf, strings.NewReader(l+"\n")) require.NoError(t, err) } mockThrottler := noopThrottler{} mockBatcher := lineBatcher{} params := internal.WriteParams{ OrgName: "my-org", BucketName: "my-bucket", Precision: api.WRITEPRECISION_US, } cli := internal.CLI{ActiveConfig: config.Config{Org: "my-default-org"}} ctrl := gomock.NewController(t) client := mock.NewMockWriteApi(ctrl) var writtenLines []string client.EXPECT().PostWrite(gomock.Any()).Return(api.ApiPostWriteRequest{ApiService: client}).Times(len(inLines)) client.EXPECT().PostWriteExecute(tmock.MatchedBy(func(in api.ApiPostWriteRequest) bool { return assert.Equal(t, params.OrgName, *in.GetOrg()) && assert.Equal(t, params.BucketName, *in.GetBucket()) && assert.Equal(t, params.Precision, *in.GetPrecision()) && assert.Equal(t, "gzip", *in.GetContentEncoding()) // Make sure the body is properly marked for compression. })).DoAndReturn(func(in api.ApiPostWriteRequest) error { writtenLines = append(writtenLines, string(in.GetBody())) return nil }).Times(len(inLines)) clients := internal.WriteClients{ Reader: &mockReader, Throttler: &mockThrottler, Writer: &mockBatcher, Client: client, } require.NoError(t, cli.Write(context.Background(), &clients, ¶ms)) require.Equal(t, inLines, writtenLines) require.True(t, mockThrottler.used) } func TestWriteOrgFromConfig(t *testing.T) { t.Parallel() inLines := []string{"fake line protocol 1", "fake line protocol 2", "fake line protocol 3"} mockReader := bufferReader{} for _, l := range inLines { _, err := io.Copy(&mockReader.buf, strings.NewReader(l+"\n")) require.NoError(t, err) } mockThrottler := noopThrottler{} mockBatcher := lineBatcher{} params := internal.WriteParams{ BucketName: "my-bucket", Precision: api.WRITEPRECISION_US, } cli := internal.CLI{ActiveConfig: config.Config{Org: "my-default-org"}} ctrl := gomock.NewController(t) client := mock.NewMockWriteApi(ctrl) var writtenLines []string client.EXPECT().PostWrite(gomock.Any()).Return(api.ApiPostWriteRequest{ApiService: client}).Times(len(inLines)) client.EXPECT().PostWriteExecute(tmock.MatchedBy(func(in api.ApiPostWriteRequest) bool { return assert.Equal(t, cli.ActiveConfig.Org, *in.GetOrg()) && assert.Equal(t, params.BucketName, *in.GetBucket()) && assert.Equal(t, params.Precision, *in.GetPrecision()) && assert.Equal(t, "gzip", *in.GetContentEncoding()) // Make sure the body is properly marked for compression. })).DoAndReturn(func(in api.ApiPostWriteRequest) error { writtenLines = append(writtenLines, string(in.GetBody())) return nil }).Times(len(inLines)) clients := internal.WriteClients{ Reader: &mockReader, Throttler: &mockThrottler, Writer: &mockBatcher, Client: client, } require.NoError(t, cli.Write(context.Background(), &clients, ¶ms)) require.Equal(t, inLines, writtenLines) require.True(t, mockThrottler.used) } func TestWriteDryRun(t *testing.T) { t.Parallel() inLines := ` fake line protocol 1 fake line protocol 2 fake line protocol 3 ` mockReader := bufferReader{} _, err := io.Copy(&mockReader.buf, strings.NewReader(inLines)) require.NoError(t, err) ctrl := gomock.NewController(t) stdio := mock.NewMockStdIO(ctrl) bytesWritten := bytes.Buffer{} stdio.EXPECT().Write(gomock.Any()).DoAndReturn(bytesWritten.Write).AnyTimes() cli := internal.CLI{ActiveConfig: config.Config{Org: "my-default-org"}, StdIO: stdio} require.NoError(t, cli.WriteDryRun(context.Background(), &mockReader)) require.Equal(t, inLines, bytesWritten.String()) }