go/ds/testutil/testutil.go

package testutil

import (
	"context"
	"fmt"
	"math/rand"
	"net"
	"net/http"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"go.skia.org/infra/go/ds"
	"go.skia.org/infra/go/emulators"
	"go.skia.org/infra/go/emulators/gcp_emulator"
	"go.skia.org/infra/go/util"
	"google.golang.org/api/iterator"
)

func cleanup(t require.TestingT, kinds ...ds.Kind) {
	for _, kind := range kinds {
		q := ds.NewQuery(kind).KeysOnly()
		it := ds.DS.Run(context.TODO(), q)
		for {
			k, err := it.Next(nil)
			if err == iterator.Done {
				break
			} else if err != nil {
				t.Errorf("Failed to clean database: %s", err)
				t.FailNow()
			}
			err = ds.DS.Delete(context.TODO(), k)
			require.NoError(t, err)
		}
	}
}

// InitDatastore is a common utility function used in tests. It sets up the
// datastore to connect to the emulator and also clears out all instances of
// the given 'kinds' from the datastore.
func InitDatastore(t require.TestingT, kinds ...ds.Kind) util.CleanupFunc {
	gcp_emulator.RequireDatastore(t.(*testing.T))
	// Copied from net/http to create a fresh http client. In some tests the
	// httpmock replaces the default http client and the healthcheck below fails.
	var transport http.RoundTripper = &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: (&net.Dialer{
			Timeout:   30 * time.Second,
			KeepAlive: 30 * time.Second,
			DualStack: true,
		}).DialContext,
		MaxIdleConns:          100,
		IdleConnTimeout:       90 * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
	}
	httpClient := &http.Client{Transport: transport}

	// Do a quick healthcheck against the host, which will fail immediately if it's down.
	emulatorHost := emulators.GetEmulatorHostEnvVar(emulators.Datastore)
	_, err := httpClient.Get("http://" + emulatorHost + "/")
	require.NoError(t, err, fmt.Sprintf("Cloud emulator host %s appears to be down or not accessible.", emulatorHost))

	r := rand.New(rand.NewSource(time.Now().UnixNano()))
	err = ds.InitForTesting("test-project", fmt.Sprintf("test-namespace-%d", r.Uint64()))
	require.NoError(t, err)
	cleanup(t, kinds...)
	return func() {
		cleanup(t, kinds...)
	}
}