Add container reset feature + postgres impl

docker.go

@@ -119,6 +119,15 @@ func (d *docker) startContainer(ctx context.Context, image string, ports NamedPo
 
 		if ok {
 			d.log.Info("re-using container")
+
+			if cfg.resetFunc != nil {
+				d.log.Info("resetting container")
+
+				if err := cfg.resetFunc(container); err != nil {
+					return nil, fmt.Errorf("can't reset container state: %w", err)
+				}
+			}
+
 			return container, nil
 		}
 	}

go.mod

@@ -71,7 +71,6 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/sirupsen/logrus v1.8.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	github.com/stretchr/objx v0.4.0 // indirect
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
 	github.com/xdg-go/scram v1.1.1 // indirect
 	github.com/xdg-go/stringprep v1.0.3 // indirect

go.sum

@@ -44,8 +44,6 @@ github.com/PuerkitoBio/purell v1.1.1 h1:WEQqlqaGbrPkxLJWfBwQmfEAE1Z7ONdDLqrN38tN
 github.com/PuerkitoBio/purell v1.1.1/go.mod h1:c11w/QuzBsJSee3cPx9rAFu61PvFxuPbtSwDGJws/X0=
 github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 h1:d+Bc7a5rLufV/sSk/8dngufqelfh6jnri85riMAaF/M=
 github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578/go.mod h1:uGdkoq3SwY9Y+13GIhn11/XLaGBb4BfwItxLd5jeuXE=
-github.com/aws/aws-sdk-go v1.44.105 h1:UUwoD1PRKIj3ltrDUYTDQj5fOTK3XsnqolLpRTMmSEM=
-github.com/aws/aws-sdk-go v1.44.105/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo=
 github.com/aws/aws-sdk-go v1.44.119 h1:TPkpDsanBMcZaF5wHwpKhjkapRV/b7d2qdC+a+IPbmY=
 github.com/aws/aws-sdk-go v1.44.119/go.mod h1:y4AeaBuwd2Lk+GepC1E9v0qOiTws0MIWAX4oIKwKHZo=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
@@ -72,15 +70,11 @@ github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.0-20210816181553-5444fa50b93d/go.
 github.com/decred/dcrd/dcrec/secp256k1/v4 v4.0.1/go.mod h1:hyedUtir6IdtD/7lIxGeCxkaw7y45JueMRL4DIyJDKs=
 github.com/deepmap/oapi-codegen v1.10.1 h1:xybuJUR6D8l7P+LAuxOm5SD7nTlFKHWvOPl31q+DDVs=
 github.com/deepmap/oapi-codegen v1.10.1/go.mod h1:TvVmDQlUkFli9gFij/gtW1o+tFBr4qCHyv2zG+R0YZY=
-github.com/denisenkom/go-mssqldb v0.12.2 h1:1OcPn5GBIobjWNd+8yjfHNIaFX14B1pWI3F9HZy5KXw=
-github.com/denisenkom/go-mssqldb v0.12.2/go.mod h1:lnIw1mZukFRZDJYQ0Pb833QS2IaC3l5HkEfra2LJ+sk=
 github.com/denisenkom/go-mssqldb v0.12.3 h1:pBSGx9Tq67pBOTLmxNuirNTeB8Vjmf886Kx+8Y+8shw=
 github.com/denisenkom/go-mssqldb v0.12.3/go.mod h1:k0mtMFOnU+AihqFxPMiF05rtiDrorD1Vrm1KEz5hxDo=
 github.com/dnaeon/go-vcr v1.2.0/go.mod h1:R4UdLID7HZT3taECzJs4YgbbH6PIGXB6W/sc5OLb6RQ=
 github.com/docker/distribution v2.8.0+incompatible h1:l9EaZDICImO1ngI+uTifW+ZYvvz7fKISBAKpg+MbWbY=
 github.com/docker/distribution v2.8.0+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w=
-github.com/docker/docker v20.10.18+incompatible h1:SN84VYXTBNGn92T/QwIRPlum9zfemfitN7pbsp26WSc=
-github.com/docker/docker v20.10.18+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/docker v20.10.20+incompatible h1:kH9tx6XO+359d+iAkumyKDc5Q1kOwPuAUaeri48nD6E=
 github.com/docker/docker v20.10.20+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.4.0 h1:El9xVISelRB7BuFusrZozjnkIM5YnzCViNKohAFqRJQ=
@@ -216,8 +210,6 @@ github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpO
 github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
 github.com/imdario/mergo v0.3.12 h1:b6R2BslTbIEToALKP7LxUvijTsNI9TAe80pLWN2g/HU=
 github.com/imdario/mergo v0.3.12/go.mod h1:jmQim1M+e3UYxmgPu/WyfjB3N3VflVyUjjjwH0dnCYA=
-github.com/influxdata/influxdb-client-go/v2 v2.10.0 h1:bWCwNsp0KxBioW9PTG7LPk7/uXj2auHezuUMpztbpZY=
-github.com/influxdata/influxdb-client-go/v2 v2.10.0/go.mod h1:x7Jo5UHHl+w8wu8UnGiNobDDHygojXwJX4mx7rXGKMk=
 github.com/influxdata/influxdb-client-go/v2 v2.11.0 h1:BrHYv38rWkAnp22gIaHFp5LpOCazOqRMRvVE1yW3ym8=
 github.com/influxdata/influxdb-client-go/v2 v2.11.0/go.mod h1:YteV91FiQxRdccyJ2cHvj2f/5sq4y4Njqu1fQzsQCOU=
 github.com/influxdata/line-protocol v0.0.0-20210922203350-b1ad95c89adf h1:7JTmneyiNEwVBOHSjoMxiWAqB992atOeepeFYegn5RU=
@@ -331,7 +323,6 @@ github.com/stoewer/go-strcase v1.2.0/go.mod h1:IBiWB2sKIp3wVVQ3Y035++gc+knqhUQag
 github.com/streadway/amqp v1.0.0 h1:kuuDrUJFZL1QYL9hUNuCxNObNzB0bV/ZG5jV3RWAQgo=
 github.com/streadway/amqp v1.0.0/go.mod h1:AZpEONHx3DKn8O/DFsRAY58/XVQiIPMTMB1SddzLXVw=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0 h1:M2gUjqZET1qApGOWNSnZ49BAIMX4F/1plDv3+l31EJ4=
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
@@ -368,8 +359,6 @@ github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9de
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
-go.mongodb.org/mongo-driver v1.10.2 h1:4Wk3cnqOrQCn0P92L3/mmurMxzdvWWs5J9jinAVKD+k=
-go.mongodb.org/mongo-driver v1.10.2/go.mod h1:z4XpeoU6w+9Vht+jAFyLgVrD+jGSQQe0+CBWFHNiHt8=
 go.mongodb.org/mongo-driver v1.10.3 h1:XDQEvmh6z1EUsXuIkXE9TaVeqHw6SwS1uf93jFs0HBA=
 go.mongodb.org/mongo-driver v1.10.3/go.mod h1:z4XpeoU6w+9Vht+jAFyLgVrD+jGSQQe0+CBWFHNiHt8=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
@@ -489,8 +478,6 @@ golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20220923202941-7f9b1623fab7 h1:ZrnxWX62AgTKOSagEqxvb3ffipvEDX2pl7E1TdqLqIc=
-golang.org/x/sync v0.0.0-20220923202941-7f9b1623fab7/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
@@ -726,16 +713,10 @@ honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
-k8s.io/api v0.25.2 h1:v6G8RyFcwf0HR5jQGIAYlvtRNrxMJQG1xJzaSeVnIS8=
-k8s.io/api v0.25.2/go.mod h1:qP1Rn4sCVFwx/xIhe+we2cwBLTXNcheRyYXwajonhy0=
 k8s.io/api v0.25.3 h1:Q1v5UFfYe87vi5H7NU0p4RXC26PPMT8KOpr1TLQbCMQ=
 k8s.io/api v0.25.3/go.mod h1:o42gKscFrEVjHdQnyRenACrMtbuJsVdP+WVjqejfzmI=
-k8s.io/apimachinery v0.25.2 h1:WbxfAjCx+AeN8Ilp9joWnyJ6xu9OMeS/fsfjK/5zaQs=
-k8s.io/apimachinery v0.25.2/go.mod h1:hqqA1X0bsgsxI6dXsJ4HnNTBOmJNxyPp8dw3u2fSHwA=
 k8s.io/apimachinery v0.25.3 h1:7o9ium4uyUOM76t6aunP0nZuex7gDf8VGwkR5RcJnQc=
 k8s.io/apimachinery v0.25.3/go.mod h1:jaF9C/iPNM1FuLl7Zuy5b9v+n35HGSh6AQ4HYRkCqwo=
-k8s.io/client-go v0.25.2 h1:SUPp9p5CwM0yXGQrwYurw9LWz+YtMwhWd0GqOsSiefo=
-k8s.io/client-go v0.25.2/go.mod h1:i7cNU7N+yGQmJkewcRD2+Vuj4iz7b30kI8OcL3horQ4=
 k8s.io/client-go v0.25.3 h1:oB4Dyl8d6UbfDHD8Bv8evKylzs3BXzzufLiO27xuPs0=
 k8s.io/client-go v0.25.3/go.mod h1:t39LPczAIMwycjcXkVc+CB+PZV69jQuNx4um5ORDjQA=
 k8s.io/klog/v2 v2.0.0/go.mod h1:PBfzABfn139FHAV07az/IF9Wp1bkk3vpT2XSJ76fSDE=

options.go

@@ -197,6 +197,19 @@ func WithContainerReuse() Option {
 	}
 }
 
+// WithContainerReset can be used together with WithContainerReuse to perform
+// arbitrary work on the reusable container. For example, this work might
+// include deleting data created by other tests, or restoring the desired
+// container state in any other way.
+//
+// Initialization function is called on the container after the cleanup
+// completes.
+func WithContainerReset(f func(*Container) error) Option {
+	return func(o *Options) {
+		o.resetFunc = f
+	}
+}
+
 // HealthcheckFunc defines a function to be used to determine container health.
 // It receives a host and a port, and returns an error if the container is not
 // ready, or nil when the container can be used. One example of HealthcheckFunc
@@ -288,6 +301,7 @@ type Options struct {
 	healthcheck         HealthcheckFunc
 	healthcheckInterval time.Duration
 	logWriter           io.Writer
+	resetFunc           ResetFunc
 }
 
 func buildConfig(opts ...Option) *Options {

preset.go

@@ -19,3 +19,11 @@ type Preset interface {
 	// Preset implementation.
 	Options() []Option
 }
+
+// ResetFunc is a function that can be called to reset container state between
+// test provided to the WithContainerReset option.
+//
+// Some presets might implement a generic ResetFunc that should be fine for
+// most use cases. For custom usage, it can be implemented depending on a
+// use-case.
+type ResetFunc func(*Container) error

preset/postgres/preset.go

@@ -21,6 +21,19 @@ const (
 	defaultVersion  = "12.5"
 )
 
+var (
+	persistentSchemas = map[string]any{
+		"pg_toast":           nil,
+		"pg_temp_1":          nil,
+		"pg_toast_temp_1":    nil,
+		"pg_catalog":         nil,
+		"information_schema": nil,
+	}
+	persistentUsers = map[string]any{
+		"postgres": nil,
+	}
+)
+
 func init() {
 	registry.Register("postgres", func() gnomock.Preset { return &P{} })
 }
@@ -173,3 +186,80 @@ func connect(c *gnomock.Container, db string) (*sql.DB, error) {
 
 	return conn, conn.Ping()
 }
+
+func Reset() gnomock.ResetFunc {
+	return func(c *gnomock.Container) error {
+		db, err := connect(c, defaultDatabase)
+		if err != nil {
+			return fmt.Errorf("database connection failed: %w", err)
+		}
+
+		defer func() { _ = db.Close() }()
+
+		if err := dropSchemas(db); err != nil {
+			return fmt.Errorf("can't drop schemas: %w", err)
+		}
+
+		if err := dropUsers(db); err != nil {
+			return fmt.Errorf("can't drop users: %w", err)
+		}
+
+		return nil
+	}
+}
+
+func dropSchemas(db *sql.DB) error {
+	rows, err := db.Query("select nspname from pg_catalog.pg_namespace")
+	if err != nil {
+		return fmt.Errorf("can't list postgres namespaces: %w", err)
+	}
+
+	defer func() { _ = rows.Close() }()
+
+	for rows.Next() {
+		var namespace string
+
+		if err := rows.Scan(&namespace); err != nil {
+			return fmt.Errorf("unexpected namespace value: %w", err)
+		}
+
+		if _, ok := persistentSchemas[namespace]; !ok {
+			dropStmt := fmt.Sprintf("drop schema %s cascade", namespace)
+			if _, err := db.Exec(dropStmt); err != nil {
+				return fmt.Errorf("can't drop namespace %s: %w", namespace, err)
+			}
+		}
+	}
+
+	if _, err := db.Exec("create schema public"); err != nil {
+		return fmt.Errorf("can't create public schema: %w", err)
+	}
+
+	return nil
+}
+
+func dropUsers(db *sql.DB) error {
+	rows, err := db.Query("select usename from pg_catalog.pg_user")
+	if err != nil {
+		return fmt.Errorf("can't list postgres users: %w", err)
+	}
+
+	defer func() { _ = rows.Close() }()
+
+	for rows.Next() {
+		var user string
+
+		if err := rows.Scan(&user); err != nil {
+			return fmt.Errorf("unexpected user value: %w", err)
+		}
+
+		if _, ok := persistentUsers[user]; !ok {
+			dropStmt := fmt.Sprintf("drop user %s", user)
+			if _, err := db.Exec(dropStmt); err != nil {
+				return fmt.Errorf("can't drop role %s: %w", user, err)
+			}
+		}
+	}
+
+	return nil
+}

preset/postgres/reset_test.go

@@ -0,0 +1,75 @@
+package postgres_test
+
+import (
+	"database/sql"
+	"fmt"
+	"testing"
+
+	"github.com/orlangure/gnomock"
+	"github.com/orlangure/gnomock/preset/postgres"
+	"github.com/stretchr/testify/suite"
+)
+
+func TestReset(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, new(ResetTestSuite))
+}
+
+type ResetTestSuite struct {
+	suite.Suite
+
+	container *gnomock.Container
+	connStr   string
+}
+
+func (t *ResetTestSuite) TearDownSuite() {
+	t.NoError(gnomock.Stop(t.container))
+}
+
+func (t *ResetTestSuite) SetupTest() {
+	p := postgres.Preset(
+		postgres.WithQueries(
+			"create table foo(a text)",
+			"insert into foo(a) values ('b')",
+		),
+		postgres.WithUser("gnomock", "foobar"),
+	)
+
+	c, err := gnomock.Start(
+		p,
+		gnomock.WithContainerName("postgres-reuse"),
+		gnomock.WithContainerReuse(),
+		gnomock.WithContainerReset(postgres.Reset()),
+	)
+	t.NoError(err)
+	t.NotNil(c)
+
+	t.container = c
+	t.connStr = fmt.Sprintf(
+		"host=%s port=%d user=%s password=%s  dbname=%s sslmode=disable",
+		c.Host, c.DefaultPort(),
+		"gnomock", "foobar", "postgres",
+	)
+}
+
+func (t *ResetTestSuite) TestResetFirst() {
+	db, err := sql.Open("postgres", t.connStr)
+	t.NoError(err)
+	t.NoError(db.Ping())
+	t.T().Cleanup(func() { t.NoError(db.Close()) })
+
+	_, err = db.Exec("insert into foo(a) values ('c'), ('d')")
+	t.NoError(err)
+}
+
+func (t *ResetTestSuite) TestResetSecond() {
+	db, err := sql.Open("postgres", t.connStr)
+	t.NoError(err)
+	t.NoError(db.Ping())
+	t.T().Cleanup(func() { t.NoError(db.Close()) })
+
+	var value string
+	err = db.QueryRow("select * from foo").Scan(&value)
+	t.NoError(err)
+	t.Equal("b", value)
+}