kratos-bootstrap/database/influxdb/client_test.go

package influxdb

import (
	"context"
	"testing"
	"time"

	"github.com/InfluxCommunity/influxdb3-go/v2/influxdb3"
	"github.com/go-kratos/kratos/v2/log"
	"github.com/stretchr/testify/assert"
	"github.com/tx7do/go-utils/trans"
	conf "github.com/tx7do/kratos-bootstrap/api/gen/go/conf/v1"
	"google.golang.org/protobuf/types/known/timestamppb"
)

type Candle struct {
	Symbol    *string
	Open      *float64
	High      *float64
	Low       *float64
	Close     *float64
	Volume    *float64
	StartTime *timestamppb.Timestamp
}

func (c *Candle) GetSymbol() string {
	if c.Symbol != nil {
		return *c.Symbol
	}
	return ""
}

func (c *Candle) GetOpen() float64 {
	if c.Open != nil {
		return *c.Open
	}
	return 0.0
}

func (c *Candle) GetHigh() float64 {
	if c.High != nil {
		return *c.High
	}
	return 0.0
}

func (c *Candle) GetLow() float64 {
	if c.Low != nil {
		return *c.Low
	}
	return 0.0
}

func (c *Candle) GetClose() float64 {
	if c.Close != nil {
		return *c.Close
	}
	return 0.0
}

func (c *Candle) GetVolume() float64 {
	if c.Volume != nil {
		return *c.Volume
	}
	return 0.0
}

func (c *Candle) GetStartTime() *timestamppb.Timestamp {
	if c.StartTime != nil {
		return c.StartTime
	}
	return timestamppb.Now()
}

type CandleMapper struct{}

var candleMapper CandleMapper

func (m *CandleMapper) ToPoint(data *Candle) *influxdb3.Point {
	p := influxdb3.NewPoint(
		"candles",
		map[string]string{"s": data.GetSymbol()},
		nil,
		data.StartTime.AsTime(),
	)

	p.
		SetDoubleField("o", data.GetOpen()).
		SetDoubleField("h", data.GetHigh()).
		SetDoubleField("l", data.GetLow()).
		SetDoubleField("c", data.GetClose()).
		SetDoubleField("v", data.GetVolume())

	return p
}

func (m *CandleMapper) ToData(point *influxdb3.Point) *Candle {
	symbol, _ := point.GetTag("s")

	return &Candle{
		Symbol:    &symbol,
		Open:      point.GetDoubleField("o"),
		High:      point.GetDoubleField("h"),
		Low:       point.GetDoubleField("l"),
		Close:     point.GetDoubleField("c"),
		Volume:    point.GetDoubleField("v"),
		StartTime: timestamppb.New(point.Values.Timestamp),
	}
}

func createTestClient() *Client {
	cli, _ := NewClient(
		log.DefaultLogger,
		&conf.Bootstrap{
			Data: &conf.Data{
				Influxdb: &conf.Data_InfluxDB{
					Host:         "http://localhost:8181",
					Token:        "apiv3_yYde4mJo0BYC7Ipi_00ZEex-A8if4swdqTBXiO-lCUDKhsIavHlRCQfo3p_DzI7S34ADHOC7Qxf600VVgW6LQQ",
					Database:     "finances",
					Organization: "primary",
				},
			},
		},
	)
	return cli
}

func TestNewClient(t *testing.T) {
	client := createTestClient()
	assert.NotNil(t, client)
}

func TestClient_Insert(t *testing.T) {
	client := createTestClient()
	assert.NotNil(t, client)

	item := &Candle{
		StartTime: timestamppb.New(time.Now()),
		Symbol:    trans.Ptr("AAPL"),
		Open:      trans.Ptr(1.0),
		High:      trans.Ptr(2.0),
		Low:       trans.Ptr(3.0),
		Close:     trans.Ptr(4.0),
		Volume:    trans.Ptr(1000.0),
	}

	point := candleMapper.ToPoint(item)

	err := client.Insert(context.Background(), point)
	assert.NoError(t, err)
}

func TestClient_BatchInsert(t *testing.T) {
	client := createTestClient()
	assert.NotNil(t, client)

	items := []*Candle{
		{
			StartTime: timestamppb.New(time.Now()),
			Symbol:    trans.Ptr("AAPL"),
			Open:      trans.Ptr(1.0),
			High:      trans.Ptr(2.0),
			Low:       trans.Ptr(3.0),
			Close:     trans.Ptr(4.0),
			Volume:    trans.Ptr(1000.0),
		},
	}

	var points []*influxdb3.Point
	for _, item := range items {
		point := candleMapper.ToPoint(item)
		points = append(points, point)
	}

	err := client.BatchInsert(
		context.Background(),
		points,
	)
	assert.NoError(t, err)
}

func TestClient_Query(t *testing.T) {
	client := createTestClient()
	assert.NotNil(t, client)

	ctx := context.Background()

	sql := `SELECT * FROM candles`

	iterator, err := client.Query(ctx, sql)
	assert.NoError(t, err)

	for iterator.Next() {
		point, _ := iterator.AsPoints().AsPoint()
		candle := candleMapper.ToData(point)
		t.Logf("[%v] Candle: %s, Open: %f, High: %f, Low: %f, Close: %f, Volume: %f\n",
			candle.GetStartTime().AsTime().String(),
			candle.GetSymbol(),
			candle.GetOpen(), candle.GetHigh(), candle.GetLow(), candle.GetClose(), candle.GetVolume(),
		)
	}

	candles, err := Query(ctx, client, sql, &candleMapper)
	assert.NoError(t, err)
	for _, candle := range candles {
		t.Logf("Candle: %s, Open: %f, High: %f, Low: %f, Close: %f, Volume: %f\n",
			candle.GetSymbol(),
			candle.GetOpen(), candle.GetHigh(), candle.GetLow(), candle.GetClose(), candle.GetVolume(),
		)
	}
}