From 281a09fe4e2d61825edec8583530caa602c2814e Mon Sep 17 00:00:00 2001
From: Michael Golowka <72365+pcman312@users.noreply.github.com>
Date: Tue, 15 Sep 2020 12:08:27 -0600
Subject: [PATCH] DBPW - Add test helpers for DB v5 interface (#9948)

---
 .../newdbplugin/testing/test_helpers.go       | 85 +++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 sdk/database/newdbplugin/testing/test_helpers.go

diff --git a/sdk/database/newdbplugin/testing/test_helpers.go b/sdk/database/newdbplugin/testing/test_helpers.go
new file mode 100644
index 0000000000..05e8d8f9b7
--- /dev/null
+++ b/sdk/database/newdbplugin/testing/test_helpers.go
@@ -0,0 +1,85 @@
+package dbtesting
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/hashicorp/vault/sdk/database/newdbplugin"
+)
+
+func getRequestTimeout(t *testing.T) time.Duration {
+	rawDur := os.Getenv("VAULT_TEST_DATABASE_REQUEST_TIMEOUT")
+	if rawDur == "" {
+		return 2 * time.Second
+	}
+
+	dur, err := time.ParseDuration(rawDur)
+	if err != nil {
+		t.Fatalf("Failed to parse custom request timeout %q: %s", rawDur, err)
+	}
+	return dur
+}
+
+func AssertInitialize(t *testing.T, db newdbplugin.Database, req newdbplugin.InitializeRequest) newdbplugin.InitializeResponse {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), getRequestTimeout(t))
+	defer cancel()
+
+	resp, err := db.Initialize(ctx, req)
+	if err != nil {
+		t.Fatalf("Failed to initialize: %s", err)
+	}
+	return resp
+}
+
+func AssertNewUser(t *testing.T, db newdbplugin.Database, req newdbplugin.NewUserRequest) newdbplugin.NewUserResponse {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), getRequestTimeout(t))
+	defer cancel()
+
+	resp, err := db.NewUser(ctx, req)
+	if err != nil {
+		t.Fatalf("Failed to create new user: %s", err)
+	}
+
+	if resp.Username == "" {
+		t.Fatalf("Missing username from NewUser response")
+	}
+	return resp
+}
+
+func AssertUpdateUser(t *testing.T, db newdbplugin.Database, req newdbplugin.UpdateUserRequest) {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), getRequestTimeout(t))
+	defer cancel()
+
+	_, err := db.UpdateUser(ctx, req)
+	if err != nil {
+		t.Fatalf("Failed to update user: %s", err)
+	}
+}
+
+func AssertDeleteUser(t *testing.T, db newdbplugin.Database, req newdbplugin.DeleteUserRequest) {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), getRequestTimeout(t))
+	defer cancel()
+
+	_, err := db.DeleteUser(ctx, req)
+	if err != nil {
+		t.Fatalf("Failed to delete user %q: %s", req.Username, err)
+	}
+}
+
+func AssertClose(t *testing.T, db newdbplugin.Database) {
+	t.Helper()
+	err := db.Close()
+	if err != nil {
+		t.Fatalf("Failed to close database: %s", err)
+	}
+}