mtg/stats/prometheus_test.go

package stats_test

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"testing"
	"time"

	"github.com/9seconds/mtg/v2/events"
	"github.com/9seconds/mtg/v2/mtglib"
	"github.com/9seconds/mtg/v2/stats"
	"github.com/stretchr/testify/suite"
)

type PrometheusTestSuite struct {
	suite.Suite

	httpListener net.Listener
	factory      *stats.PrometheusFactory
	prometheus   events.Observer
}

func (suite *PrometheusTestSuite) Get() (string, error) {
	addr := fmt.Sprintf("http://%s/", suite.httpListener.Addr().String())

	resp, err := http.Get(addr) // nolint: noctx
	if err != nil {
		return "", err // nolint: wrapcheck
	}

	defer resp.Body.Close()

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err // nolint: wrapcheck
	}

	return string(data), nil
}

func (suite *PrometheusTestSuite) SetupTest() {
	suite.httpListener, _ = net.Listen("tcp", "127.0.0.1:0")
	suite.factory = stats.NewPrometheus("mtg", "/")
	suite.prometheus = suite.factory.Make()

	go suite.factory.Serve(suite.httpListener) // nolint: errcheck
}

func (suite *PrometheusTestSuite) TearDownTest() {
	suite.prometheus.Shutdown()
	suite.NoError(suite.factory.Close())
	suite.httpListener.Close()
}

func (suite *PrometheusTestSuite) TestEventStartFinish() {
	suite.prometheus.EventStart(mtglib.EventStart{
		CreatedAt: time.Now(),
		ConnID:    "connID",
		RemoteIP:  net.ParseIP("10.0.0.10"),
	})
	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_type="ipv4"} 1`)

	suite.prometheus.EventConnectedToDC(mtglib.EventConnectedToDC{
		CreatedAt: time.Now(),
		ConnID:    "connID",
		RemoteIP:  net.ParseIP("10.0.0.1"),
		DC:        4,
	})
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_telegram_connections{dc="4",ip="10.0.0.1",ip_type="ipv4"} 1`)

	suite.prometheus.EventTraffic(mtglib.EventTraffic{
		CreatedAt: time.Now(),
		ConnID:    "connID",
		Traffic:   200,
		IsRead:    true,
	})
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_traffic{dc="4",direction="client",ip="10.0.0.1",ip_type="ipv4"} 200`)

	suite.prometheus.EventTraffic(mtglib.EventTraffic{
		CreatedAt: time.Now(),
		ConnID:    "connID",
		Traffic:   100,
		IsRead:    false,
	})
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_traffic{dc="4",direction="telegram",ip="10.0.0.1",ip_type="ipv4"} 100`)

	suite.prometheus.EventFinish(mtglib.EventFinish{
		CreatedAt: time.Now(),
		ConnID:    "connID",
	})
	time.Sleep(100 * time.Millisecond)

	data, err = suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_client_connections{ip_type="ipv4"} 0`)
	suite.Contains(data, `mtg_telegram_connections{dc="4",ip="10.0.0.1",ip_type="ipv4"} 0`)
	suite.Contains(data, `mtg_traffic{dc="4",direction="client",ip="10.0.0.1",ip_type="ipv4"} 200`)
	suite.Contains(data, `mtg_traffic{dc="4",direction="telegram",ip="10.0.0.1",ip_type="ipv4"} 100`)
}

func (suite *PrometheusTestSuite) TestEventConcurrencyLimited() {
	suite.prometheus.EventConcurrencyLimited(mtglib.EventConcurrencyLimited{
		CreatedAt: time.Now(),
	})

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_concurrency_limited 1`)
}

func (suite *PrometheusTestSuite) TestEventIPBlocklisted() {
	suite.prometheus.EventIPBlocklisted(mtglib.EventIPBlocklisted{
		CreatedAt: time.Now(),
		RemoteIP:  net.ParseIP("2001:db8::68"),
	})

	time.Sleep(100 * time.Millisecond)

	data, err := suite.Get()
	suite.NoError(err)
	suite.Contains(data, `mtg_ip_blocklisted{ip_type="ipv6"} 1`)
}

func TestPrometheus(t *testing.T) {
	t.Parallel()
	suite.Run(t, &PrometheusTestSuite{})
}