test: added TestPGDatabaseTimeout to ensure our timing behaves as it should

Signed-off-by: Pierre Souchay <pierre.souchay@pelico.io>

Author: Pierre SOUCHAY

collector/collector.go

@@ -178,8 +178,6 @@ func (p PostgresCollector) Describe(ch chan<- *prometheus.Desc) {
 
 // Collect implements the prometheus.Collector interface.
 func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
-ctx, cancel := context.WithTimeout(context.Background(), p.CollectionTimeout)
-defer cancel()
 // copy the instance so that concurrent scrapes have independent instances
 inst := p.instance.copy()
 
@@ -190,6 +188,12 @@ func (p PostgresCollector) Collect(ch chan<- prometheus.Metric) {
 p.logger.Error("Error opening connection to database", "err", err)
 return
 }
+p.collectFromConnection(inst, ch)
+}
+
+func (p PostgresCollector) collectFromConnection(inst *instance, ch chan<- prometheus.Metric) {
+ctx, cancel := context.WithTimeout(context.Background(), p.CollectionTimeout)
+defer cancel()
 
 wg := sync.WaitGroup{}
 wg.Add(len(p.Collectors))

collector/pg_database_test.go

@@ -15,13 +15,80 @@ package collector
 import (
 "context"
 "testing"
+"time"
 
 "github.com/DATA-DOG/go-sqlmock"
 "github.com/prometheus/client_golang/prometheus"
+
 dto "github.com/prometheus/client_model/go"
+"github.com/prometheus/common/promslog"
 "github.com/smartystreets/goconvey/convey"
 )
 
+// We ensure that when the database respond after a long time
+// The collection process still occurs in a predictable manner
+// Will avoid accumulation of queries on a completely frozen DB
+func TestPGDatabaseTimeout(t *testing.T) {
+
+timeoutForQuery := time.Duration(100 * time.Millisecond)
+
+db, mock, err := sqlmock.New()
+if err != nil {
+t.Fatalf("Error opening a stub db connection: %s", err)
+}
+defer db.Close()
+
+inst := &instance{db: db}
+
+columns := []string{"pg_roles.rolname", "pg_roles.rolconnlimit"}
+rows := sqlmock.NewRows(columns).AddRow("role1", 2)
+mock.ExpectQuery(pgRolesConnectionLimitsQuery).
+WillDelayFor(30 * time.Second).
+WillReturnRows(rows)
+
+log_config := promslog.Config{}
+
+logger := promslog.New(&log_config)
+
+c, err := NewPostgresCollector(logger, []string{}, "postgresql://local", []string{}, CollectionTimeout(timeoutForQuery.String()))
+if err != nil {
+t.Fatalf("error creating NewPostgresCollector: %s", err)
+}
+collector_config := collectorConfig{
+logger: logger,
+excludeDatabases: []string{},
+}
+
+collector, err := NewPGRolesCollector(collector_config)
+if err != nil {
+t.Fatalf("error creating collector: %s", err)
+}
+c.Collectors["test"] = collector
+c.instance = inst
+
+ch := make(chan prometheus.Metric)
+defer close(ch)
+
+go func() {
+for {
+<-ch
+time.Sleep(1 * time.Millisecond)
+}
+}()
+
+startTime := time.Now()
+c.collectFromConnection(inst, ch)
+elapsed := time.Since(startTime)
+
+if elapsed <= timeoutForQuery {
+t.Errorf("elapsed time was %v, should be bigger than timeout=%v", elapsed, timeoutForQuery)
+}
+
+if err := mock.ExpectationsWereMet(); err != nil {
+t.Errorf("there were unfulfilled exceptions: %s", err)
+}
+}
+
 func TestPGDatabaseCollector(t *testing.T) {
 db, mock, err := sqlmock.New()
 if err != nil {