package testutils import ( "context" "fmt" "slices" "strings" "testing" "time" "github.com/samber/lo" "github.com/openmeterio/openmeter/openmeter/meter" "github.com/openmeterio/openmeter/openmeter/streaming" "github.com/openmeterio/openmeter/pkg/filter" ) var _ streaming.Connector = &MockStreamingConnector{} func NewMockStreamingConnector(t testing.TB) *MockStreamingConnector { t.Helper() out := &MockStreamingConnector{} out.Reset() return out } type SimpleEvent struct { MeterSlug string Value float64 Time time.Time StoredAt time.Time } type MockStreamingConnector struct { rows map[string][]meter.MeterQueryRow events map[string][]SimpleEvent } func (m *MockStreamingConnector) Reset() { m.rows = map[string][]meter.MeterQueryRow{} m.events = map[string][]SimpleEvent{} } type AddOption func(event *SimpleEvent) func WithStoredAt(storedAt time.Time) AddOption { return func(event *SimpleEvent) { event.StoredAt = storedAt } } func (m *MockStreamingConnector) AddSimpleEvent(meterSlug string, value float64, at time.Time, opts ...AddOption) { event := SimpleEvent{ MeterSlug: meterSlug, Value: value, Time: at, StoredAt: at, } for _, opt := range opts { opt(&event) } m.events[meterSlug] = append(m.events[meterSlug], event) m.sortMeterEvents(meterSlug) } func (m *MockStreamingConnector) SetSimpleEvents(meterSlug string, fn func(events []SimpleEvent) []SimpleEvent) { if _, ok := m.events[meterSlug]; !ok { m.events[meterSlug] = []SimpleEvent{} } m.events[meterSlug] = fn(m.events[meterSlug]) m.sortMeterEvents(meterSlug) } func (m *MockStreamingConnector) AddRow(meterSlug string, row meter.MeterQueryRow) { m.rows[meterSlug] = append(m.rows[meterSlug], row) } func (m *MockStreamingConnector) sortMeterEvents(meterSlug string) { // Let's sort events by Time ASC slices.SortStableFunc(m.events[meterSlug], func(a, b SimpleEvent) int { return a.Time.Compare(b.Time) }) } func (c *MockStreamingConnector) CreateNamespace(ctx context.Context, namespace string) error { return nil } func (c *MockStreamingConnector) DeleteNamespace(ctx context.Context, namespace string) error { return nil } func (m *MockStreamingConnector) CountEvents(ctx context.Context, namespace string, params streaming.CountEventsParams) ([]streaming.CountEventRow, error) { return []streaming.CountEventRow{}, nil } func (m *MockStreamingConnector) ListEvents(ctx context.Context, namespace string, params streaming.ListEventsParams) ([]streaming.RawEvent, error) { return []streaming.RawEvent{}, nil } func (m *MockStreamingConnector) ListEventsV2(ctx context.Context, params streaming.ListEventsV2Params) ([]streaming.RawEvent, error) { return []streaming.RawEvent{}, nil } // Returns the result query set for the given params. If the query set is not found, // it will try to approximate the result by aggregating the simple events func (m *MockStreamingConnector) QueryMeter(ctx context.Context, namespace string, mm meter.Meter, params streaming.QueryParams) ([]meter.MeterQueryRow, error) { rows := []meter.MeterQueryRow{} _, rowOk := m.rows[mm.Key] if rowOk { for _, row := range m.rows[mm.Key] { if row.WindowStart.Equal(*params.From) && row.WindowEnd.Equal(*params.To) { rows = append(rows, row) } } } else { row, err := m.aggregateEvents(mm, params) if err != nil { return rows, err } rows = append(rows, row...) } return rows, nil } func (m *MockStreamingConnector) BatchInsert(ctx context.Context, events []streaming.RawEvent) error { return nil } func (m *MockStreamingConnector) ValidateJSONPath(ctx context.Context, jsonPath string) (bool, error) { return strings.HasPrefix(jsonPath, "$."), nil } func (m *MockStreamingConnector) windowSizeDuration(windowSize meter.WindowSize) time.Duration { switch windowSize { case meter.WindowSizeMinute: return time.Minute case meter.WindowSizeHour: return time.Hour case meter.WindowSizeDay: return 24 * time.Hour default: return 0 } } // filterStoredAt evaluates a FilterTimeUnix predicate against storedAt using Unix-second precision, // matching the ClickHouse stored_at column behavior. Composite $and/$or filters are evaluated // recursively. func filterStoredAt(f *filter.FilterTimeUnix, storedAt time.Time) bool { if f == nil || f.IsEmpty() { return true } unix := storedAt.Unix() switch { case f.Gt != nil: return unix > f.Gt.Unix() case f.Gte != nil: return unix >= f.Gte.Unix() case f.Lt != nil: return unix < f.Lt.Unix() case f.Lte != nil: return unix <= f.Lte.Unix() case f.And != nil: for _, sub := range *f.And { if !filterStoredAt(&filter.FilterTimeUnix{FilterTime: sub}, storedAt) { return false } } return true case f.Or != nil: for _, sub := range *f.Or { if filterStoredAt(&filter.FilterTimeUnix{FilterTime: sub}, storedAt) { return true } } return false default: return true } } // We approximate the actual logic by a simple filter + aggregation for most cases func (m *MockStreamingConnector) aggregateEvents(mm meter.Meter, params streaming.QueryParams) ([]meter.MeterQueryRow, error) { events, ok := m.events[mm.Key] if !ok { return []meter.MeterQueryRow{}, meter.NewMeterNotFoundError(mm.Key) } if params.From == nil || params.To == nil { return nil, fmt.Errorf("streaming mock connector does not support filtering without from and to") } if params.FilterStoredAt != nil && !params.FilterStoredAt.IsEmpty() { events = lo.Filter(events, func(event SimpleEvent, _ int) bool { return filterStoredAt(params.FilterStoredAt, event.StoredAt) }) } // Let's truncate the window size to the second, as clickhouse does not support sub-second precision from := params.From.Truncate(streaming.MinimumWindowSizeDuration) to := params.To.Truncate(streaming.MinimumWindowSizeDuration) rows := make([]meter.MeterQueryRow, 0) if params.WindowSize != nil && params.WindowTimeZone != nil { // TODO: windowtimezone will be ignored windowingStart, _ := params.WindowSize.Truncate(from) // The first truncated time that from query falls into windowingEnd, _ := params.WindowSize.Truncate(to) // The last truncated time that to query falls into if !to.Equal(windowingEnd) { windowingEnd, _ = params.WindowSize.AddTo(windowingEnd) } numOfWindows := int(windowingEnd.Sub(windowingStart).Seconds()) / int(m.windowSizeDuration(*params.WindowSize).Seconds()) if numOfWindows == 0 { return nil, fmt.Errorf("couldnt calculate windows") } for i := 0; i < numOfWindows; i++ { rows = append(rows, meter.MeterQueryRow{ Value: 0, WindowStart: windowingStart.Add(m.windowSizeDuration(*params.WindowSize) * time.Duration(i)), WindowEnd: windowingStart.Add(m.windowSizeDuration(*params.WindowSize) * time.Duration(i+1)), GroupBy: map[string]*string{}, }) } } else { rows = append(rows, meter.MeterQueryRow{ Value: 0, WindowStart: from, WindowEnd: to, GroupBy: map[string]*string{}, }) } for i := range rows { row := &rows[i] var value float64 effectiveWindowSize := lo.FromPtrOr(params.WindowSize, streaming.MinimumWindowSize) for _, event := range events { eventWindowStart, err := effectiveWindowSize.Truncate(event.Time) if err != nil { return nil, fmt.Errorf("failed to truncate by windowsize in event aggregation") } // windowend is exclusive when doing this rounding eventWindowEnd, err := effectiveWindowSize.AddTo(eventWindowStart) if err != nil { return nil, fmt.Errorf("failed calculate window end in event aggregation") } if (eventWindowStart.After(row.WindowStart) || eventWindowStart.Equal(row.WindowStart)) && (eventWindowEnd.Before(row.WindowEnd) || eventWindowEnd.Equal(row.WindowEnd)) { // TODO: Add support for more aggregation types switch mm.Aggregation { case meter.MeterAggregationLatest: // Note: events are already sorted by time ASC when they are registered value = event.Value default: value += event.Value } } } rows[i].Value = value } // Clickhouse doesn't return tumpled result rows if there are no rows (events) in the tumpled period // To simulate this for the SUM behavior, we simply filter out rows that have 0 value rows = lo.Filter(rows, func(row meter.MeterQueryRow, _ int) bool { return row.Value != 0 }) return rows, nil } func (m *MockStreamingConnector) ListSubjects(ctx context.Context, params streaming.ListSubjectsParams) ([]string, error) { return []string{}, nil } func (m *MockStreamingConnector) ListGroupByValues(ctx context.Context, params streaming.ListGroupByValuesParams) ([]string, error) { return []string{}, nil }