63 lines
1.4 KiB
Go
63 lines
1.4 KiB
Go
package mock
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
|
|
"github.com/influxdata/influxdb/v2"
|
|
"github.com/influxdata/influxdb/v2/models"
|
|
)
|
|
|
|
// PointsWriter is a mock structure for writing points.
|
|
type PointsWriter struct {
|
|
timesWriteCalled int
|
|
mu sync.RWMutex
|
|
Points []models.Point
|
|
Err error
|
|
|
|
WritePointsFn func(ctx context.Context, orgID influxdb.ID, bucketID influxdb.ID, points []models.Point) error
|
|
}
|
|
|
|
// ForceError is for error testing, if WritePoints is called after ForceError, it will return that error.
|
|
func (p *PointsWriter) ForceError(err error) {
|
|
p.mu.Lock()
|
|
p.Err = err
|
|
p.mu.Unlock()
|
|
}
|
|
|
|
// WritePoints writes points to the PointsWriter that will be exposed in the Values.
|
|
func (p *PointsWriter) WritePoints(ctx context.Context, orgID influxdb.ID, bucketID influxdb.ID, points []models.Point) error {
|
|
if p.WritePointsFn != nil {
|
|
return p.WritePointsFn(ctx, orgID, bucketID, points)
|
|
}
|
|
|
|
p.mu.Lock()
|
|
p.timesWriteCalled++
|
|
p.Points = append(p.Points, points...)
|
|
err := p.Err
|
|
p.mu.Unlock()
|
|
return err
|
|
}
|
|
|
|
// Next returns the next (oldest) batch of values.
|
|
func (p *PointsWriter) Next() models.Point {
|
|
var points models.Point
|
|
p.mu.RLock()
|
|
if len(p.Points) == 0 {
|
|
p.mu.RUnlock()
|
|
return points
|
|
}
|
|
p.mu.RUnlock()
|
|
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
points, p.Points = p.Points[0], p.Points[1:]
|
|
return points
|
|
}
|
|
|
|
func (p *PointsWriter) WritePointsCalled() int {
|
|
p.mu.Lock()
|
|
defer p.mu.Unlock()
|
|
return p.timesWriteCalled
|
|
}
|