Wrapper around gcloud datastore emulator -- for use in tests...

emulator_test.go

package datastore_test

import (
	"bufio"
	"bytes"
	"context"
	"fmt"
	"net/http"
	"os"
	"os/exec"
	"strings"
	"time"
)

// NewEmulator returns a new instance of Emulator.
func NewEmulator() (*Emulator, error) {
	e := &Emulator{}
	if err := e.Start(); err != nil {
		return nil, err
	}
	return e, nil
}

// Emulator manages the GCP Datastore Emulator process.
type Emulator struct {
	Host      string
	ProjectID string
}

// Start starts the emulator which involves initializing the environment,
// starting the emulator and blocking until correct startup is confirmed.
func (e *Emulator) Start() error {
	if err := e.initEnv(); err != nil {
		return err
	}
	if err := e.command(
		"start",
		"--consistency=1.0",  // prevents random test failures
		"--no-store-on-disk", // test in memory
	).Start(); err != nil {
		return err
	}
	if err := e.confirmStartup(); err != nil {
		return err
	}
	return nil
}

// Reset resets the Datastore Emulator (but only works in testing/i.e. when
// using in-memory storage).
func (e *Emulator) Reset() error {
	return e.request("/reset", http.MethodPost)
}

// Close terminates the emulator process and cleans up the environemental
// variables.
func (e *Emulator) Close() error {
	os.Unsetenv("DATASTORE_EMULATOR_HOST")
	os.Unsetenv("DATASTORE_PROJECT_ID")
	if e.healthCheck() {
		return e.request("/shutdown", http.MethodPost)
	}
	return nil
}

func (e *Emulator) initEnv() error {
	var buf bytes.Buffer
	cmd := e.command("env-init")
	cmd.Stdout = &buf
	if err := cmd.Run(); err != nil {
		return err
	}
	env := make(map[string]string, 5)
	scanner := bufio.NewScanner(&buf)
	for scanner.Scan() {
		line := scanner.Text()
		e := strings.Split(strings.Split(line, " ")[1], "=")
		env[e[0]] = e[1]
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	e.Host = env["DATASTORE_HOST"]
	e.ProjectID = env["DATASTORE_PROJECT_ID"]
	os.Setenv("DATASTORE_EMULATOR_HOST", env["DATASTORE_EMULATOR_HOST"])
	os.Setenv("DATASTORE_PROJECT_ID", env["DATASTORE_PROJECT_ID"])
	return nil
}

func (e *Emulator) healthCheck() bool {
	if err := e.request("", http.MethodGet); err != nil {
		return false
	}
	return true
}

func (e *Emulator) confirmStartup() error {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	t := time.NewTicker(200 * time.Millisecond)
	for {
		select {
		case <-t.C:
			if e.healthCheck() {
				t.Stop()
				return nil
			}
		case <-ctx.Done():
			t.Stop()
			return ctx.Err()
		}
	}
}

func (e *Emulator) command(extraArgs ...string) *exec.Cmd {
	args := []string{"beta", "emulators", "datastore"}
	args = append(args, extraArgs...)
	return exec.Command("gcloud", args...)
}

func (e *Emulator) request(path, method string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, method, e.Host+path, nil)
	if err != nil {
		return err
	}
	c := http.Client{}
	resp, err := c.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode != 200 {
		return fmt.Errorf("status code error: %d", resp.StatusCode)
	}
	return nil
}