Merge pull request #2511 from jhaals/setupDatabase

Add tests for database setup

cmd/yopass-server/main.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/jhaals/yopass/pkg/server"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/collectors"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/spf13/pflag"
 	"github.com/spf13/viper"
@@ -47,28 +48,13 @@ func init() {
 
 func main() {
 	logger := configureZapLogger()
-
-	var db server.Database
-	switch database := viper.GetString("database"); database {
-	case "memcached":
-		memcached := viper.GetString("memcached")
-		db = server.NewMemcached(memcached)
-		logger.Debug("configured Memcached", zap.String("address", memcached))
-	case "redis":
-		redis := viper.GetString("redis")
-		var err error
-		db, err = server.NewRedis(redis)
-		if err != nil {
-			logger.Fatal("invalid Redis URL", zap.Error(err))
-		}
-		logger.Debug("configured Redis", zap.String("url", redis))
-	default:
-		logger.Fatal("unsupported database, expected 'memcached' or 'redis'", zap.String("database", database))
+	db, err := setupDatabase(logger)
+	if err != nil {
+		logger.Fatal("failed to setup database", zap.Error(err))
 	}
-
 	registry := prometheus.NewRegistry()
-	registry.MustRegister(prometheus.NewProcessCollector(prometheus.ProcessCollectorOpts{}))
-	registry.MustRegister(prometheus.NewGoCollector())
+	registry.MustRegister(collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}))
+	registry.MustRegister(collectors.NewGoCollector())
 
 	cert := viper.GetString("tls-cert")
 	key := viper.GetString("tls-key")
@@ -146,3 +132,24 @@ func configureZapLogger() *zap.Logger {
 	zap.ReplaceGlobals(logger)
 	return logger
 }
+
+func setupDatabase(logger *zap.Logger) (server.Database, error) {
+	var db server.Database
+	switch database := viper.GetString("database"); database {
+	case "memcached":
+		memcached := viper.GetString("memcached")
+		db = server.NewMemcached(memcached)
+		logger.Debug("configured Memcached", zap.String("address", memcached))
+	case "redis":
+		redis := viper.GetString("redis")
+		var err error
+		db, err = server.NewRedis(redis)
+		if err != nil {
+			return nil, fmt.Errorf("invalid Redis URL: %w", err)
+		}
+		logger.Debug("configured Redis", zap.String("url", redis))
+	default:
+		return nil, fmt.Errorf("unsupported database, expected 'memcached' or 'redis' got '%s'", database)
+	}
+	return db, nil
+}

cmd/yopass-server/main_test.go

@@ -0,0 +1,70 @@
+package main
+
+import (
+	"testing"
+
+	"github.com/spf13/viper"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+	"go.uber.org/zap/zaptest/observer"
+)
+
+func TestSetupDatabase(t *testing.T) {
+	viper.Set("database", "memcached")
+	viper.Set("memcached", "localhost:11211")
+
+	core, logs := observer.New(zapcore.DebugLevel)
+	logger := zap.New(core)
+
+	db, _ := setupDatabase(logger)
+	if db == nil {
+		t.Fatal("Expected non-nil database")
+	}
+
+	if logs.FilterMessage("configured Memcached").Len() != 1 {
+		t.Error("Expected log message 'configured Memcached'")
+	}
+}
+
+func TestSetupDatabaseRedis(t *testing.T) {
+	viper.Set("database", "redis")
+	viper.Set("redis", "redis://localhost:6379/0")
+
+	core, logs := observer.New(zapcore.DebugLevel)
+	logger := zap.New(core)
+
+	db, _ := setupDatabase(logger)
+	if db == nil {
+		t.Fatal("Expected non-nil database")
+	}
+
+	if logs.FilterMessage("configured Redis").Len() != 1 {
+		t.Error("Expected log message 'configured Redis'")
+	}
+}
+
+func TestSetupDatabaseRedisWithInvalidUrl(t *testing.T) {
+	viper.Set("database", "redis")
+	viper.Set("redis", "boop")
+
+	core, _ := observer.New(zapcore.DebugLevel)
+	logger := zap.New(core)
+
+	_, err := setupDatabase(logger)
+	expected := `invalid Redis URL: invalid redis URL scheme: `
+	if err.Error() != expected {
+		t.Fatalf("Expected '%s', got '%v'", expected, err.Error())
+	}
+}
+
+func TestSetupDatabaseInvalid(t *testing.T) {
+	viper.Set("database", "invalid")
+	core, _ := observer.New(zapcore.DebugLevel)
+	logger := zap.New(core)
+
+	_, err := setupDatabase(logger)
+	expected := `unsupported database, expected 'memcached' or 'redis' got 'invalid'`
+	if err.Error() != expected {
+		t.Fatalf("Expected '%s', got '%v'", expected, err.Error())
+	}
+}