github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/clients/pkg/promtail/targets/cloudflare/util_test.go (about)

     1  package cloudflare
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"time"
     7  
     8  	"github.com/cloudflare/cloudflare-go"
     9  	"github.com/stretchr/testify/mock"
    10  )
    11  
    12  type fakeCloudflareClient struct {
    13  	mock.Mock
    14  }
    15  
    16  func (f *fakeCloudflareClient) CallCount() int {
    17  	var actualCalls int
    18  	for _, call := range f.Calls {
    19  		if call.Method == "LogpullReceived" {
    20  			actualCalls++
    21  		}
    22  	}
    23  	return actualCalls
    24  }
    25  
    26  type fakeLogIterator struct {
    27  	logs    []string
    28  	current string
    29  
    30  	err error
    31  }
    32  
    33  func (f *fakeLogIterator) Next() bool {
    34  	if len(f.logs) == 0 {
    35  		return false
    36  	}
    37  	f.current = f.logs[0]
    38  	if f.current == `error` {
    39  		f.err = errors.New("error")
    40  		return false
    41  	}
    42  	f.logs = f.logs[1:]
    43  	return true
    44  }
    45  func (f *fakeLogIterator) Err() error                         { return f.err }
    46  func (f *fakeLogIterator) Line() []byte                       { return []byte(f.current) }
    47  func (f *fakeLogIterator) Fields() (map[string]string, error) { return nil, nil }
    48  func (f *fakeLogIterator) Close() error                       { return nil }
    49  
    50  func newFakeCloudflareClient() *fakeCloudflareClient {
    51  	return &fakeCloudflareClient{}
    52  }
    53  
    54  func (f *fakeCloudflareClient) LogpullReceived(ctx context.Context, start, end time.Time) (cloudflare.LogpullReceivedIterator, error) {
    55  	r := f.Called(ctx, start, end)
    56  	if r.Get(0) != nil {
    57  		return r.Get(0).(cloudflare.LogpullReceivedIterator), nil
    58  	}
    59  	return nil, r.Error(1)
    60  }