signoz/pkg/query-service/tests/integration/test_utils.go

package tests

import (
	"fmt"
	"testing"
	"time"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/google/uuid"
	"github.com/jmoiron/sqlx"
	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/entry"
	mockhouse "github.com/srikanthccv/ClickHouse-go-mock"
	"github.com/stretchr/testify/require"
	"go.signoz.io/signoz/pkg/query-service/app/clickhouseReader"
	"go.signoz.io/signoz/pkg/query-service/interfaces"
	"go.signoz.io/signoz/pkg/query-service/model"
	"golang.org/x/exp/maps"
)

func NewMockClickhouseReader(
	t *testing.T, testDB *sqlx.DB, featureFlags interfaces.FeatureLookup,
) (
	*clickhouseReader.ClickHouseReader, mockhouse.ClickConnMockCommon,
) {
	require.NotNil(t, testDB)

	mockDB, err := mockhouse.NewClickHouseWithQueryMatcher(nil, sqlmock.QueryMatcherRegexp)

	require.Nil(t, err, "could not init mock clickhouse")
	reader := clickhouseReader.NewReaderFromClickhouseConnection(
		mockDB,
		clickhouseReader.NewOptions("", 10, 10, 10*time.Second, ""),
		testDB,
		"",
		featureFlags,
		"",
	)

	return reader, mockDB
}

func addLogsQueryExpectation(
	mockClickhouse mockhouse.ClickConnMockCommon,
	logsToReturn []model.SignozLog,
) {
	cols := []mockhouse.ColumnType{}
	cols = append(cols, mockhouse.ColumnType{Type: "UInt64", Name: "timestamp"})
	cols = append(cols, mockhouse.ColumnType{Type: "UInt64", Name: "observed_timestamp"})
	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "id"})
	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "trace_id"})
	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "span_id"})
	cols = append(cols, mockhouse.ColumnType{Type: "UInt32", Name: "trace_flags"})
	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "severity_text"})
	cols = append(cols, mockhouse.ColumnType{Type: "UInt8", Name: "severity_number"})
	cols = append(cols, mockhouse.ColumnType{Type: "String", Name: "body"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "resources_string_key"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "resources_string_value"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "attributes_string_key"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "attributes_string_value"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "attributes_int64_key"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(Int64)", Name: "attributes_int64_value"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "attributes_float64_key"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(Float64)", Name: "attributes_float64_value"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(String)", Name: "attributes_bool_key"})
	cols = append(cols, mockhouse.ColumnType{Type: "Array(Bool)", Name: "attributes_bool_value"})

	values := [][]any{}
	for _, l := range logsToReturn {
		rowValues := []any{}
		rowValues = append(rowValues, l.Timestamp)
		rowValues = append(rowValues, l.Timestamp)
		rowValues = append(rowValues, l.ID)
		rowValues = append(rowValues, l.TraceID)
		rowValues = append(rowValues, l.SpanID)
		rowValues = append(rowValues, l.TraceFlags)
		rowValues = append(rowValues, l.SeverityText)
		rowValues = append(rowValues, l.SeverityNumber)
		rowValues = append(rowValues, l.Body)
		rowValues = append(rowValues, maps.Keys(l.Resources_string))
		rowValues = append(rowValues, maps.Values(l.Resources_string))
		rowValues = append(rowValues, maps.Keys(l.Attributes_string))
		rowValues = append(rowValues, maps.Values(l.Attributes_string))
		rowValues = append(rowValues, maps.Keys(l.Attributes_int64))
		rowValues = append(rowValues, maps.Values(l.Attributes_int64))
		rowValues = append(rowValues, maps.Keys(l.Attributes_float64))
		rowValues = append(rowValues, maps.Values(l.Attributes_float64))
		rowValues = append(rowValues, maps.Keys(l.Attributes_bool))
		rowValues = append(rowValues, maps.Values(l.Attributes_bool))
		values = append(values, rowValues)
	}

	rows := mockhouse.NewRows(cols, values)
	mockClickhouse.ExpectQuery(
		"SELECT .*? from signoz_logs.distributed_logs.*",
	).WillReturnRows(rows)
}

func makeTestSignozLog(
	body string,
	attributes map[string]interface{},
) model.SignozLog {

	testLog := model.SignozLog{
		Timestamp:          uint64(time.Now().UnixNano()),
		Body:               body,
		Attributes_bool:    map[string]bool{},
		Attributes_string:  map[string]string{},
		Attributes_int64:   map[string]int64{},
		Attributes_float64: map[string]float64{},
		Resources_string:   map[string]string{},
		SeverityText:       entry.Info.String(),
		SeverityNumber:     uint8(entry.Info),
		SpanID:             uuid.New().String(),
		TraceID:            uuid.New().String(),
	}

	for k, v := range attributes {
		switch v.(type) {
		case bool:
			testLog.Attributes_bool[k] = v.(bool)
		case string:
			testLog.Attributes_string[k] = v.(string)
		case int:
			testLog.Attributes_int64[k] = int64(v.(int))
		case float64:
			testLog.Attributes_float64[k] = v.(float64)
		default:
			panic(fmt.Sprintf("found attribute value of unsupported type %T in test log", v))
		}
	}

	return testLog
}