Refactoring and adding more tests (#130)

* Adding tests

* Revert rename

* Cleanup

* Fix lint

* More lint fixes

Author: davissp14

internal/api/handle_admin.go

@@ -189,7 +189,7 @@ func handleUpdatePostgresSettings(w http.ResponseWriter, r *http.Request) {
 		}}
 	}
 
-	err = flypg.PushUserConfig(node.PGConfig, consul)
+	err = flypg.PushUserConfig(&node.PGConfig, consul)
 	if err != nil {
 		renderErr(w, err)
 		return
@@ -218,7 +218,7 @@ func handleApplyConfig(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	err = flypg.SyncUserConfig(node.PGConfig, consul)
+	err = flypg.SyncUserConfig(&node.PGConfig, consul)
 	if err != nil {
 		renderErr(w, err)
 		return

internal/flypg/admin/admin.go

@@ -9,6 +9,11 @@ import (
 	"github.com/jackc/pgx/v5"
 )
 
+type Credential struct {
+	Username string
+	Password string
+}
+
 func GrantAccess(ctx context.Context, pg *pgx.Conn, username string) error {
 	sql := fmt.Sprintf("GRANT pg_read_all_data, pg_write_all_data TO %q", username)
 	_, err := pg.Exec(ctx, sql)
@@ -28,6 +33,38 @@ func CreateUser(ctx context.Context, pg *pgx.Conn, username string, password str
 	return err
 }
 
+func ManageDefaultUsers(ctx context.Context, conn *pgx.Conn, creds []Credential) error {
+	curUsers, err := ListUsers(ctx, conn)
+	if err != nil {
+		return fmt.Errorf("failed to list existing users: %s", err)
+	}
+
+	for _, c := range creds {
+		exists := false
+		for _, curUser := range curUsers {
+			if c.Username == curUser.Username {
+				exists = true
+			}
+		}
+
+		if exists {
+			if err := ChangePassword(ctx, conn, c.Username, c.Password); err != nil {
+				return fmt.Errorf("failed to update credentials for user %s: %s", c.Username, err)
+			}
+		} else {
+			if err := CreateUser(ctx, conn, c.Username, c.Password); err != nil {
+				return fmt.Errorf("failed to create require user %s: %s", c.Username, err)
+			}
+
+			if err := GrantSuperuser(ctx, conn, c.Username); err != nil {
+				return fmt.Errorf("failed to grant superuser privileges to user %s: %s", c.Username, err)
+			}
+		}
+	}
+
+	return nil
+}
+
 func ChangePassword(ctx context.Context, pg *pgx.Conn, username, password string) error {
 	sql := fmt.Sprintf("ALTER USER %s WITH LOGIN PASSWORD '%s';", username, password)
 

internal/flypg/config.go

@@ -52,6 +52,10 @@ func PushUserConfig(c Config, consul *state.Store) error {
 }
 
 func SyncUserConfig(c Config, consul *state.Store) error {
+	if os.Getenv("UNIT_TESTING") != "" {
+		return nil
+	}
+
 	cfg, err := pullFromConsul(c, consul)
 	if err != nil {
 		return fmt.Errorf("failed to pull config from consul: %s", err)

internal/flypg/flypg.go

@@ -2,7 +2,6 @@ package flypg
 
 import (
 	"fmt"
-	"os"
 	"time"
 )
 
@@ -12,25 +11,8 @@ type FlyPGConfig struct {
 
 	internalConfig ConfigMap
 	userConfig     ConfigMap
-
-	configPath string
 }
 
-func (c *FlyPGConfig) SetDefaults() {
-	c.internalConfig = ConfigMap{
-		"deadMemberRemovalThreshold": time.Hour * 24,
-	}
-}
-
-func NewInternalConfig(configPath string) *FlyPGConfig {
-	return &FlyPGConfig{
-		internalConfigFilePath: fmt.Sprintf("%s/flypg.internal.conf", configPath),
-		userConfigFilePath:     fmt.Sprintf("%s/flypg.user.conf", configPath),
-		configPath:             configPath,
-		internalConfig:         ConfigMap{},
-		userConfig:             ConfigMap{},
-	}
-}
 func (c *FlyPGConfig) InternalConfig() ConfigMap {
 	return c.internalConfig
 }
@@ -55,6 +37,12 @@ func (c *FlyPGConfig) UserConfigFile() string {
 	return c.userConfigFilePath
 }
 
+func (c *FlyPGConfig) SetDefaults() {
+	c.internalConfig = ConfigMap{
+		"deadMemberRemovalThreshold": time.Hour * 24,
+	}
+}
+
 func (c *FlyPGConfig) CurrentConfig() (ConfigMap, error) {
 	internal, err := ReadFromFile(c.InternalConfigFile())
 	if err != nil {
@@ -80,28 +68,11 @@ func (c *FlyPGConfig) CurrentConfig() (ConfigMap, error) {
 func (c *FlyPGConfig) initialize() error {
 	c.SetDefaults()
 
-	file, err := os.Create(c.internalConfigFilePath)
-	if err != nil {
-		return err
-	}
-	defer func() { _ = file.Close() }()
-
-	if err := file.Sync(); err != nil {
-		return fmt.Errorf("failed to sync file: %s", err)
-	} else if err := file.Close(); err != nil {
-		return fmt.Errorf("failed to close file: %s", err)
-	}
-
-	file, err = os.Create(c.userConfigFilePath)
-	if err != nil {
-		return err
-	}
-	defer func() { _ = file.Close() }()
+	// Note - Sync from consul has been disabled for this component.
+	// It will be re-enabled once we offer user-defined configuration.
 
-	if err := file.Sync(); err != nil {
-		return fmt.Errorf("failed to sync file: %s", err)
-	} else if err := file.Close(); err != nil {
-		return fmt.Errorf("failed to close file: %s", err)
+	if err := WriteConfigFiles(c); err != nil {
+		return fmt.Errorf("failed to write internal config files: %s", err)
 	}
 
 	return nil

internal/flypg/flypg_test.go

@@ -0,0 +1,51 @@
+package flypg
+
+import (
+	"testing"
+	"time"
+
+	"github.com/fly-apps/postgres-flex/internal/utils"
+)
+
+const (
+	flyTestDirectory          = "./test_results"
+	flyInternalConfigFilePath = "./test_results/flypg.internal.conf"
+	flyUserConfigFilePath     = "./test_results/flypg.user.conf"
+)
+
+func TestFlyConfigInitialization(t *testing.T) {
+	if err := setup(t); err != nil {
+		t.Fatal(err)
+	}
+	defer cleanup()
+
+	cfg := FlyPGConfig{
+		internalConfigFilePath: flyInternalConfigFilePath,
+		userConfigFilePath:     flyInternalConfigFilePath,
+	}
+
+	if err := cfg.initialize(); err != nil {
+		t.Fatal(err)
+	}
+
+	t.Run("configFiles", func(t *testing.T) {
+		if !utils.FileExists(cfg.internalConfigFilePath) {
+			t.Fatalf("expected %s to exist, but doesn't", cfg.internalConfigFilePath)
+		}
+
+		if !utils.FileExists(cfg.userConfigFilePath) {
+			t.Fatalf("expected %s to exist, but doesn't", cfg.userConfigFilePath)
+		}
+	})
+
+	t.Run("defaults", func(t *testing.T) {
+		cfg, err := cfg.CurrentConfig()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if cfg["deadMemberRemovalThreshold"] != (24 * time.Hour).String() {
+			t.Fatalf("expected deadMemberRemovalThreshold to be 24h, but got %v", cfg["deadMemberRemovalThreshold"])
+		}
+	})
+}