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

     1  package fake
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/grafana/loki/clients/pkg/promtail/api"
     7  )
     8  
     9  // Client is a fake client used for testing.
    10  type Client struct {
    11  	entries  chan api.Entry
    12  	received []api.Entry
    13  	once     sync.Once
    14  	mtx      sync.Mutex
    15  	wg       sync.WaitGroup
    16  	OnStop   func()
    17  }
    18  
    19  func New(stop func()) *Client {
    20  	c := &Client{
    21  		OnStop:  stop,
    22  		entries: make(chan api.Entry),
    23  	}
    24  	c.wg.Add(1)
    25  	go func() {
    26  		defer c.wg.Done()
    27  		for e := range c.entries {
    28  			c.mtx.Lock()
    29  			c.received = append(c.received, e)
    30  			c.mtx.Unlock()
    31  		}
    32  	}()
    33  	return c
    34  }
    35  
    36  // Stop implements client.Client
    37  func (c *Client) Stop() {
    38  	c.once.Do(func() { close(c.entries) })
    39  	c.wg.Wait()
    40  	c.OnStop()
    41  }
    42  
    43  func (c *Client) Chan() chan<- api.Entry {
    44  	return c.entries
    45  }
    46  
    47  func (c *Client) Received() []api.Entry {
    48  	c.mtx.Lock()
    49  	defer c.mtx.Unlock()
    50  	cpy := make([]api.Entry, len(c.received))
    51  	copy(cpy, c.received)
    52  	return cpy
    53  }
    54  
    55  // StopNow implements client.Client
    56  func (c *Client) StopNow() {
    57  	c.Stop()
    58  }
    59  
    60  func (c *Client) Name() string {
    61  	return "fake"
    62  }
    63  
    64  // Clear is used to cleanup the buffered received entries, so the same client can be re-used between
    65  // test cases.
    66  func (c *Client) Clear() {
    67  	c.mtx.Lock()
    68  	defer c.mtx.Unlock()
    69  	c.received = []api.Entry{}
    70  }