Spin out cmd tests

Author: bplunkett-stripe

cmd/pg-schema-diff/apply_cmd.go

@@ -49,7 +49,7 @@ func buildApplyCmd() *cobra.Command {
 
 		cmd.SilenceUsage = true
 
-		plan, err := generatePlan(context.Background(), generatePlanParameters{
+		plan, err := generatePlan(cmd.Context(), generatePlanParameters{
 			fromSchema:       fromSchema,
 			toSchema:         toSchema,
 			tempDbConnConfig: connConfig,
@@ -59,12 +59,12 @@ func buildApplyCmd() *cobra.Command {
 		if err != nil {
 			return err
 		} else if len(plan.Statements) == 0 {
-			fmt.Println("Schema matches expected. No plan generated")
+			cmd.Println("Schema matches expected. No plan generated")
 			return nil
 		}
 
-		fmt.Println(header("Review plan"))
-		fmt.Print(planToPrettyS(plan), "\n\n")
+		cmd.Println(header("Review plan"))
+		cmd.Print(planToPrettyS(plan), "\n\n")
 
 		if err := failIfHazardsNotAllowed(plan, *allowedHazardsTypesStrs); err != nil {
 			return err
@@ -81,10 +81,10 @@ func buildApplyCmd() *cobra.Command {
 			}
 		}
 
-		if err := runPlan(context.Background(), connConfig, plan); err != nil {
+		if err := runPlan(cmd.Context(), cmd, connConfig, plan); err != nil {
 			return err
 		}
-		fmt.Println("Schema applied successfully")
+		cmd.Println("Schema applied successfully")
 		return nil
 	}
 
@@ -123,7 +123,7 @@ func failIfHazardsNotAllowed(plan diff.Plan, allowedHazardsTypesStrs []string) e
 	return nil
 }
 
-func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) error {
+func runPlan(ctx context.Context, cmd *cobra.Command, connConfig *pgx.ConnConfig, plan diff.Plan) error {
 	connPool, err := openDbWithPgxConfig(connConfig)
 	if err != nil {
 		return err
@@ -143,8 +143,8 @@ func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) er
 	// must be executed within its own transaction block. Postgres will error if you try to set a TRANSACTION-level
 	// timeout for it. SESSION-level statement_timeouts are respected by `ADD INDEX CONCURRENTLY`
 	for i, stmt := range plan.Statements {
-		fmt.Println(header(fmt.Sprintf("Executing statement %d", getDisplayableStmtIdx(i))))
-		fmt.Printf("%s\n\n", statementToPrettyS(stmt))
+		cmd.Println(header(fmt.Sprintf("Executing statement %d", getDisplayableStmtIdx(i))))
+		cmd.Printf("%s\n\n", statementToPrettyS(stmt))
 		start := time.Now()
 		if _, err := conn.ExecContext(ctx, fmt.Sprintf("SET SESSION statement_timeout = %d", stmt.Timeout.Milliseconds())); err != nil {
 			return fmt.Errorf("setting statement timeout: %w", err)
@@ -155,9 +155,9 @@ func runPlan(ctx context.Context, connConfig *pgx.ConnConfig, plan diff.Plan) er
 		if _, err := conn.ExecContext(ctx, stmt.ToSQL()); err != nil {
 			return fmt.Errorf("executing migration statement. the database maybe be in a dirty state: %s: %w", stmt, err)
 		}
-		fmt.Printf("Finished executing statement. Duration: %s\n", time.Since(start))
+		cmd.Printf("Finished executing statement. Duration: %s\n", time.Since(start))
 	}
-	fmt.Println(header("Complete"))
+	cmd.Println(header("Complete"))
 
 	return nil
 }

cmd/pg-schema-diff/apply_cmd_test.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"github.com/stripe/pg-schema-diff/internal/pgdump"
+	"github.com/stripe/pg-schema-diff/internal/pgengine"
+)
+
+func (suite *cmdTestSuite) TestApplyCmd() {
+	// Non-comprehensive set of tests for the plan command. Not totally comprehensive to avoid needing to avoid
+	// hindering developer velocity when updating the command.
+	type testCase struct {
+		name string
+		// fromDbArg is an optional argument to override the default "--from-dsn" arg.
+		fromDbArg func(db *pgengine.DB) []string
+		args      []string
+		// dynamicArgs is function that can be used to build args that are dynamic, i.e.,
+		// saving schemas to a randomly generated temporary directory.
+		dynamicArgs []dArgGenerator
+
+		outputContains []string
+		// expectedSchema is the schema that is expected to be in the database after the migration.
+		// If nil, the expected schema will be the fromDDL.
+		expectedSchemaDDL []string
+		// expectErrContains is a list of substrings that are expected to be contained in the error returned by
+		// cmd.RunE. This is DISTINCT from stdErr.
+		expectErrContains []string
+	}
+	for _, tc := range []testCase{
+		{
+			name:        "to dir",
+			dynamicArgs: []dArgGenerator{tempSchemaDirDArg("to-dir", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name:        "to dsn",
+			dynamicArgs: []dArgGenerator{tempDsnDArg(suite.pgEngine, "to-dsn", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name: "from empty dsn",
+			fromDbArg: func(db *pgengine.DB) []string {
+				tempSetPqEnvVarsForDb(suite.T(), db)
+				return []string{"--from-empty-dsn"}
+			},
+			dynamicArgs: []dArgGenerator{tempSchemaDirDArg("to-dir", []string{"CREATE TABLE foobar();"})},
+
+			expectedSchemaDDL: []string{"CREATE TABLE foobar();"},
+		},
+		{
+			name:              "no to schema provided",
+			expectErrContains: []string{"must be set"},
+		},
+		{
+			name:              "two to schemas provided",
+			args:              []string{"--to-dir", "some-other-dir", "--to-dsn", "some-dsn"},
+			expectErrContains: []string{"only one of"},
+		},
+	} {
+		suite.Run(tc.name, func() {
+			fromDb := tempDbWithSchema(suite.T(), suite.pgEngine, nil)
+			if tc.fromDbArg == nil {
+				tc.fromDbArg = func(db *pgengine.DB) []string {
+					return []string{"--from-dsn", db.GetDSN()}
+				}
+			}
+			args := append([]string{
+				"apply",
+				"--skip-confirm-prompt",
+			}, tc.fromDbArg(fromDb)...)
+			args = append(args, tc.args...)
+			suite.runCmdWithAssertions(runCmdWithAssertionsParams{
+				args:              args,
+				dynamicArgs:       tc.dynamicArgs,
+				outputContains:    tc.outputContains,
+				expectErrContains: tc.expectErrContains,
+			})
+			// The migration should have been successful. Assert it was.
+			expectedDb := tempDbWithSchema(suite.T(), suite.pgEngine, tc.expectedSchemaDDL)
+			expectedDbDump, err := pgdump.GetDump(expectedDb, pgdump.WithSchemaOnly())
+			suite.Require().NoError(err)
+			fromDbDump, err := pgdump.GetDump(fromDb, pgdump.WithSchemaOnly())
+			suite.Require().NoError(err)
+
+			suite.Equal(expectedDbDump, fromDbDump)
+		})
+	}
+}

cmd/pg-schema-diff/main.go

@@ -6,21 +6,19 @@ import (
 	"github.com/spf13/cobra"
 )
 
-// rootCmd represents the base command when called without any subcommands
-var rootCmd = &cobra.Command{
-	Use:   "pg-schema-diff",
-	Short: "Diff two Postgres schemas and generate the SQL to get from one to the other",
-}
-
-func init() {
+func buildRootCmd() *cobra.Command {
+	rootCmd := &cobra.Command{
+		Use:   "pg-schema-diff",
+		Short: "Diff two Postgres schemas and generate the SQL to get from one to the other",
+	}
 	rootCmd.AddCommand(buildPlanCmd())
 	rootCmd.AddCommand(buildApplyCmd())
 	rootCmd.AddCommand(buildVersionCmd())
+	return rootCmd
 }
 
 func main() {
-	err := rootCmd.Execute()
-	if err != nil {
+	if err := buildRootCmd().Execute(); err != nil {
 		os.Exit(1)
 	}
 }

cmd/pg-schema-diff/main_test.go

@@ -0,0 +1,143 @@
+package main
+
+import (
+	"bytes"
+	"database/sql"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+	"github.com/stripe/pg-schema-diff/internal/pgengine"
+)
+
+type cmdTestSuite struct {
+	suite.Suite
+	pgEngine *pgengine.Engine
+}
+
+func (suite *cmdTestSuite) SetupSuite() {
+	pgEngine, err := pgengine.StartEngine()
+	suite.Require().NoError(err)
+	suite.pgEngine = pgEngine
+}
+
+func (suite *cmdTestSuite) TearDownSuite() {
+	suite.Require().NoError(suite.pgEngine.Close())
+}
+
+type runCmdWithAssertionsParams struct {
+	args []string
+	// dynamicArgs is function that can be used to build args that are dynamic, i.e.,
+	// saving schemas to a randomly generated temporary directory.
+	dynamicArgs []dArgGenerator
+
+	// outputContains is a list of substrings that are expected to be contained in the stdout output of the command.
+	outputContains []string
+	// expectErrContains is a list of substrings that are expected to be contained in the error returned by
+	// cmd.RunE. This is DISTINCT from stdErr.
+	expectErrContains []string
+}
+
+func (suite *cmdTestSuite) runCmdWithAssertions(tc runCmdWithAssertionsParams) {
+	args := tc.args
+	for _, da := range tc.dynamicArgs {
+		args = append(args, da(suite.T())...)
+	}
+
+	rootCmd := buildRootCmd()
+	rootCmd.SetArgs(args)
+	stdOut := &bytes.Buffer{}
+	rootCmd.SetOut(stdOut)
+	stdErr := &bytes.Buffer{}
+	rootCmd.SetErr(stdErr)
+
+	err := rootCmd.Execute()
+	if len(tc.expectErrContains) > 0 {
+		for _, e := range tc.expectErrContains {
+			suite.ErrorContains(err, e)
+		}
+	} else {
+		stdErrStr := stdErr.String()
+		suite.Require().NoError(err)
+		// Only assert the std error is empty if we don't expect an error
+		suite.Empty(stdErrStr, "expected no stderr")
+	}
+
+	stdOutStr := stdOut.String()
+	if len(tc.outputContains) > 0 {
+		for _, o := range tc.outputContains {
+			suite.Contains(stdOutStr, o)
+		}
+	}
+}
+
+// dArgGenerator generates argument at the run-time of the test case...
+// intended for resources that are not known at test start and potentially need
+// to be cleaned up.
+type dArgGenerator func(*testing.T) []string
+
+func tempSchemaDirDArg(argName string, ddl []string) dArgGenerator {
+	return func(t *testing.T) []string {
+		t.Helper()
+		return []string{"--" + argName, tempSchemaDir(t, ddl)}
+	}
+}
+
+func tempSchemaDir(t *testing.T, ddl []string) string {
+	t.Helper()
+	dir := t.TempDir()
+	for i, d := range ddl {
+		require.NoError(t, os.WriteFile(filepath.Join(dir, fmt.Sprintf("ddl_%d.sql", i)), []byte(d), 0644))
+	}
+	return dir
+}
+
+func tempDsnDArg(pgEngine *pgengine.Engine, argName string, ddl []string) dArgGenerator {
+	return func(t *testing.T) []string {
+		t.Helper()
+		db := tempDbWithSchema(t, pgEngine, ddl)
+		return []string{"--" + argName, db.GetDSN()}
+	}
+}
+
+func tempDbWithSchema(t *testing.T, pgEngine *pgengine.Engine, ddl []string) *pgengine.DB {
+	t.Helper()
+	db, err := pgEngine.CreateDatabase()
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, db.DropDB())
+	})
+	dbPool, err := sql.Open("pgx", db.GetDSN())
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, dbPool.Close())
+	}()
+	for _, d := range ddl {
+		_, err := dbPool.Exec(d)
+		require.NoError(t, err)
+	}
+	return db
+}
+
+func tempSetPqEnvVarsForDb(t *testing.T, db *pgengine.DB) {
+	t.Helper()
+	tempSetEnvVar(t, "PGHOST", db.GetConnOpts()[pgengine.ConnectionOptionHost])
+	tempSetEnvVar(t, "PGPORT", db.GetConnOpts()[pgengine.ConnectionOptionPort])
+	tempSetEnvVar(t, "PGUSER", db.GetConnOpts()[pgengine.ConnectionOptionUser])
+	tempSetEnvVar(t, "PGDATABASE", db.GetName())
+}
+
+func tempSetEnvVar(t *testing.T, k, v string) {
+	t.Helper()
+	require.NoError(t, os.Setenv(k, v))
+	t.Cleanup(func() {
+		require.NoError(t, os.Unsetenv(k))
+	})
+}
+
+func TestCmdTestSuite(t *testing.T) {
+	suite.Run(t, new(cmdTestSuite))
+}