Add unit tests to increase code quality

pkg/internal/testhelpers/containers_test.go

@@ -8,7 +8,9 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io/ioutil"
 	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/jackc/pgx/v4"
@@ -59,13 +61,32 @@ func TestMain(m *testing.M) {
 	flag.Parse()
 	ctx := context.Background()
 	if !testing.Short() && *useDocker {
-		container, err := StartPGContainer(ctx)
+		pgContainer, err := StartPGContainer(ctx)
+		if err != nil {
+			fmt.Println("Error setting up container", err)
+			os.Exit(1)
+		}
+		path, err := ioutil.TempDir("", "prom_test")
+		if err != nil {
+			fmt.Println("Error getting temp dir for Prometheus storage", err)
+			os.Exit(1)
+		}
+		err = os.Mkdir(filepath.Join(path, "wal"), 0777)
+		if err != nil {
+			fmt.Println("Error getting temp dir for Prometheus storage", err)
+			os.Exit(1)
+		}
+		promContainer, err := StartPromContainer(path, ctx)
 		if err != nil {
 			fmt.Println("Error setting up container", err)
 			os.Exit(1)
 		}
 		defer func() {
-			err := container.Terminate(ctx)
+			err := pgContainer.Terminate(ctx)
+			if err != nil {
+				panic(err)
+			}
+			err = promContainer.Terminate(ctx)
 			if err != nil {
 				panic(err)
 			}

pkg/util/election_test.go

@@ -23,7 +23,7 @@ import (
 var (
 	useDocker        = flag.Bool("use-docker", true, "start database using a docker container")
 	database         = flag.String("database", "tmp_db_timescale_migrate_test", "database to run integration tests on")
-	electionInterval = flag.Duration("election-interval", 5*time.Second, "Scheduled election interval")
+	electionInterval = flag.Duration("election-interval", 1*time.Second, "Scheduled election interval")
 )
 
 func TestRestElection(t *testing.T) {

pkg/util/util_test.go

@@ -4,12 +4,54 @@
 package util
 
 import (
+	"testing"
+	"time"
+
 	"github.com/timescale/timescale-prometheus/pkg/log"
 )
 
+const (
+	step      = 3.0
+	frequency = 5.0
+	count     = 5.0
+)
+
 func init() {
 	err := log.Init("debug")
 	if err != nil {
 		panic(err)
 	}
 }
+
+func TestThroughputCalc(t *testing.T) {
+
+	calc := NewThroughputCalc(time.Second / frequency)
+	ticker := time.NewTicker(time.Second / frequency)
+	stop := time.NewTimer(time.Second / 2)
+	factor := 1.0
+	var value float64
+	current := step * factor
+
+	calc.Start()
+
+	for range ticker.C {
+		calc.SetCurrent(current)
+		value = <-calc.Values
+		current = current + (step * factor)
+		select {
+		case <-stop.C:
+			if value != step*frequency*factor {
+				t.Errorf("Value is not %f, its %f", step*frequency*factor, value)
+			}
+
+			factor++
+
+			if factor > count {
+				return
+			}
+
+			stop.Reset(time.Second / 2)
+		default:
+		}
+	}
+}