<commit_before>package main

import (
	"flag"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/cli"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"math"
)

type DeleteCommand struct {
	Ui 				cli.Ui
	InstanceId 		string
	OlderThan 		string
	RequireAtLeast	int
	DryRun			bool
}

// descriptions for args
var deleteDscrInstanceId = "The EC2 instance from which the AMIs to be deleted were originally created."
var deleteOlderThan = "Delete AMIs older than the specified time; accepts formats like '30d' or '4h'."
var requireAtLeast = "Never delete AMIs such that fewer than this number of AMIs will remain. E.g. require at least 3 AMIs remain."
var deleteDscrDryRun = "Execute a simulated run. Lists AMIs to be deleted, but does not actually delete them."

func (c *DeleteCommand) Help() string {
	return `ec2-snapper create <args> [--help]

Create an AMI of the given EC2 instance.

Available args are:
--instance      	` + deleteDscrInstanceId + `
--older-than    	` + deleteOlderThan + `
--require-at-least      ` + requireAtLeast + `
--dry-run       	` + deleteDscrDryRun
}

func (c *DeleteCommand) Synopsis() string {
	return "Delete the specified AMIs"
}

func (c *DeleteCommand) Run(args []string) int {

	// Handle the command-line args
	cmdFlags := flag.NewFlagSet("delete", flag.ExitOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&c.InstanceId, "instance", "", deleteDscrInstanceId)
	cmdFlags.StringVar(&c.OlderThan, "older-than", "", deleteOlderThan)
	cmdFlags.IntVar(&c.RequireAtLeast, "require-at-least", 0, requireAtLeast)
	cmdFlags.BoolVar(&c.DryRun, "dry-run", false, deleteDscrDryRun)

	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Check for required command-line args
	if c.InstanceId == "" {
		c.Ui.Error("ERROR: The argument '--instance' is required.")
		return 1
	}

	if c.OlderThan == "" {
		c.Ui.Error("ERROR: The argument '--older-than' is required.")
		return 1
	}

	if c.RequireAtLeast < 0 {
		c.Ui.Error("ERROR: The argument '--require-at-least' must be a positive integer.")
		return 1
	}

	// Warn the user that this is a dry run
	if c.DryRun {
		c.Ui.Warn("WARNING: This is a dry run, and no actions will be taken, despite what any output may say!")
	}

	// Create an EC2 service object; AWS region is picked up from the "AWS_REGION" env var.
	svc := ec2.New(nil)

	// Get a list of the existing AMIs that were created for the given EC2 instance
	resp, err := svc.DescribeImages(&ec2.DescribeImagesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String("tag:ec2-snapper-instance-id"),
				Values: []*string{&c.InstanceId},
			},
		},
	})
	if err != nil && strings.Contains(err.Error(), "NoCredentialProviders") {
		c.Ui.Error("ERROR: No AWS credentials were found.  Either set the environment variables AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, or run this program on an EC2 instance that has an IAM Role with the appropriate permissions.")
		return 1
	} else if err != nil {
		panic(err)
	}
	if len(resp.Images) == 0 {
		c.Ui.Error("No AMIs were found for EC2 instance \"" + c.InstanceId + "\"")
		return 0
	}

	// Check that at least the --require-at-least number of AMIs exists
	// - Note that even if this passes, we still want to avoid deleting so many AMIs that we go below the threshold
	if len(resp.Images) <= c.RequireAtLeast {
		c.Ui.Info("NO ACTION TAKEN. There are currently " + strconv.Itoa(len(resp.Images)) + " AMIs, and --require-at-least=" + strconv.Itoa(c.RequireAtLeast) + " so no further action can be taken.")
		return 0
	}

	// Get the AWS Account ID of the current AWS account
	// We need this to do a more efficient lookup on the snapshot volumes
	awsAccountId := *resp.Images[0].OwnerID
	c.Ui.Output("==> Identified current AWS Account Id as " + awsAccountId)

	// Parse our date range
	match, _ := regexp.MatchString("^[0-9]*(h|d|m)$", c.OlderThan)
	if ! match {
		c.Ui.Error("The --older-than value of \"" + c.OlderThan + "\" is not formatted properly.  Use formats like 30d or 24h")
		return 0
	}

	var minutes float64
	var hours float64

	// We were given a time like "12h"
	if match, _ := regexp.MatchString("^[0-9]*(h)$", c.OlderThan); match {
		hours, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
	}

	// We were given a time like "15d"
	if match, _ := regexp.MatchString("^[0-9]*(d)$", c.OlderThan); match {
		hours, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
		hours *= 24
	}

	// We were given a time like "5m"
	if match, _ := regexp.MatchString("^[0-9]*(m)$", c.OlderThan); match {
		minutes, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
		hours = minutes/60
	}

	// Now filter the AMIs to only include those within our date range
	var filteredAmis[]*ec2.Image
	for i := 0; i < len(resp.Images); i++ {
		now := time.Now()
		creationDate, err := time.Parse(time.RFC3339Nano, *resp.Images[i].CreationDate)
		if err != nil {
			panic(err)
		}

		duration := now.Sub(creationDate)

		if duration.Hours() > hours {
			filteredAmis = append(filteredAmis, resp.Images[i])
		}
	}
	c.Ui.Output("==> Found " + strconv.Itoa(len(filteredAmis)) + " total AMIs for deletion.")

	if len(filteredAmis) == 0 {
		c.Ui.Error("No AMIs to delete.")
		return 0
	}

	// Get a list of every single snapshot in our account
	// (I wasn't able to find a better way to filter these, but suggestions welcome!)
	respDscrSnapshots, err := svc.DescribeSnapshots(&ec2.DescribeSnapshotsInput{
		OwnerIDs: []*string{&awsAccountId},
	})
	if err != nil {
		panic(err)
	}
	c.Ui.Output("==> Found " + strconv.Itoa(len(respDscrSnapshots.Snapshots)) + " total snapshots in this account.")

	// Compute whether we should delete fewer AMIs to adhere to our --require-at-least requirement
	var numTotalAmis = len(resp.Images)
	var numFilteredAmis = len(filteredAmis)
	var numAmisToRemainAfterDelete = numTotalAmis - numFilteredAmis
	var numAmisToRemoveFromFiltered = math.Max(0.0, float64(c.RequireAtLeast - numAmisToRemainAfterDelete))

	if numAmisToRemoveFromFiltered > 0.0 {
		c.Ui.Output("==> Only deleting " + strconv.Itoa(len(filteredAmis) - int(numAmisToRemoveFromFiltered)) + " total AMIs to honor '--require-at-least=" + strconv.Itoa(c.RequireAtLeast) + "'.")
	}

	// Begin deleting AMIs...
	for i := 0; i < len(filteredAmis) - int(numAmisToRemoveFromFiltered); i++ {
		// Step 1: De-register the AMI
		c.Ui.Output(*filteredAmis[i].ImageID + ": De-registering AMI named \"" + *filteredAmis[i].Name + "\"...")
		_, err := svc.DeregisterImage(&ec2.DeregisterImageInput{
			DryRun: &c.DryRun,
			ImageID: filteredAmis[i].ImageID,
		})
		if err != nil {
			if ! strings.Contains(err.Error(), "DryRunOperation") {
				panic(err)
			}
		}

		// Step 2: Delete the corresponding AMI snapshot
		// Look at the "description" for each Snapshot to see if it contains our AMI id
		snapshotId := ""
		for j := 0; j < len(respDscrSnapshots.Snapshots); j++ {
			if strings.Contains(*respDscrSnapshots.Snapshots[j].Description, *filteredAmis[i].ImageID) {
				snapshotId = *respDscrSnapshots.Snapshots[j].SnapshotID
				break
			}
		}

		c.Ui.Output(*filteredAmis[i].ImageID + ": Deleting snapshot " + snapshotId + "...")
		svc.DeleteSnapshot(&ec2.DeleteSnapshotInput{
			DryRun: &c.DryRun,
			SnapshotID: &snapshotId,
		})

		c.Ui.Output(*filteredAmis[i].ImageID + ": Done!")
		c.Ui.Output("")
	}

	if c.DryRun {
		c.Ui.Info("==> DRY RUN. Had this not been a dry run, " + strconv.Itoa(len(filteredAmis)) + " AMI's and their corresponding snapshots would have been deleted.")
	} else {
		c.Ui.Info("==> Success! Deleted " + strconv.Itoa(len(filteredAmis) - int(numAmisToRemoveFromFiltered)) + " AMI's and their corresponding snapshots.")
	}
	return 0
}

<commit_msg>Delete all snapshots, not just the first one.<commit_after>package main

import (
	"flag"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/mitchellh/cli"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"math"
)

type DeleteCommand struct {
	Ui 				cli.Ui
	InstanceId 		string
	OlderThan 		string
	RequireAtLeast	int
	DryRun			bool
}

// descriptions for args
var deleteDscrInstanceId = "The EC2 instance from which the AMIs to be deleted were originally created."
var deleteOlderThan = "Delete AMIs older than the specified time; accepts formats like '30d' or '4h'."
var requireAtLeast = "Never delete AMIs such that fewer than this number of AMIs will remain. E.g. require at least 3 AMIs remain."
var deleteDscrDryRun = "Execute a simulated run. Lists AMIs to be deleted, but does not actually delete them."

func (c *DeleteCommand) Help() string {
	return `ec2-snapper create <args> [--help]

Create an AMI of the given EC2 instance.

Available args are:
--instance      	` + deleteDscrInstanceId + `
--older-than    	` + deleteOlderThan + `
--require-at-least      ` + requireAtLeast + `
--dry-run       	` + deleteDscrDryRun
}

func (c *DeleteCommand) Synopsis() string {
	return "Delete the specified AMIs"
}

func (c *DeleteCommand) Run(args []string) int {

	// Handle the command-line args
	cmdFlags := flag.NewFlagSet("delete", flag.ExitOnError)
	cmdFlags.Usage = func() { c.Ui.Output(c.Help()) }

	cmdFlags.StringVar(&c.InstanceId, "instance", "", deleteDscrInstanceId)
	cmdFlags.StringVar(&c.OlderThan, "older-than", "", deleteOlderThan)
	cmdFlags.IntVar(&c.RequireAtLeast, "require-at-least", 0, requireAtLeast)
	cmdFlags.BoolVar(&c.DryRun, "dry-run", false, deleteDscrDryRun)

	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	// Check for required command-line args
	if c.InstanceId == "" {
		c.Ui.Error("ERROR: The argument '--instance' is required.")
		return 1
	}

	if c.OlderThan == "" {
		c.Ui.Error("ERROR: The argument '--older-than' is required.")
		return 1
	}

	if c.RequireAtLeast < 0 {
		c.Ui.Error("ERROR: The argument '--require-at-least' must be a positive integer.")
		return 1
	}

	// Warn the user that this is a dry run
	if c.DryRun {
		c.Ui.Warn("WARNING: This is a dry run, and no actions will be taken, despite what any output may say!")
	}

	// Create an EC2 service object; AWS region is picked up from the "AWS_REGION" env var.
	svc := ec2.New(nil)

	// Get a list of the existing AMIs that were created for the given EC2 instance
	resp, err := svc.DescribeImages(&ec2.DescribeImagesInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name: aws.String("tag:ec2-snapper-instance-id"),
				Values: []*string{&c.InstanceId},
			},
		},
	})
	if err != nil && strings.Contains(err.Error(), "NoCredentialProviders") {
		c.Ui.Error("ERROR: No AWS credentials were found.  Either set the environment variables AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, or run this program on an EC2 instance that has an IAM Role with the appropriate permissions.")
		return 1
	} else if err != nil {
		panic(err)
	}
	if len(resp.Images) == 0 {
		c.Ui.Error("No AMIs were found for EC2 instance \"" + c.InstanceId + "\"")
		return 0
	}

	// Check that at least the --require-at-least number of AMIs exists
	// - Note that even if this passes, we still want to avoid deleting so many AMIs that we go below the threshold
	if len(resp.Images) <= c.RequireAtLeast {
		c.Ui.Info("NO ACTION TAKEN. There are currently " + strconv.Itoa(len(resp.Images)) + " AMIs, and --require-at-least=" + strconv.Itoa(c.RequireAtLeast) + " so no further action can be taken.")
		return 0
	}

	// Get the AWS Account ID of the current AWS account
	// We need this to do a more efficient lookup on the snapshot volumes
	awsAccountId := *resp.Images[0].OwnerID
	c.Ui.Output("==> Identified current AWS Account Id as " + awsAccountId)

	// Parse our date range
	match, _ := regexp.MatchString("^[0-9]*(h|d|m)$", c.OlderThan)
	if ! match {
		c.Ui.Error("The --older-than value of \"" + c.OlderThan + "\" is not formatted properly.  Use formats like 30d or 24h")
		return 0
	}

	var minutes float64
	var hours float64

	// We were given a time like "12h"
	if match, _ := regexp.MatchString("^[0-9]*(h)$", c.OlderThan); match {
		hours, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
	}

	// We were given a time like "15d"
	if match, _ := regexp.MatchString("^[0-9]*(d)$", c.OlderThan); match {
		hours, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
		hours *= 24
	}

	// We were given a time like "5m"
	if match, _ := regexp.MatchString("^[0-9]*(m)$", c.OlderThan); match {
		minutes, _ = strconv.ParseFloat(c.OlderThan[0:len(c.OlderThan)-1], 64)
		hours = minutes/60
	}

	// Now filter the AMIs to only include those within our date range
	var filteredAmis[]*ec2.Image
	for i := 0; i < len(resp.Images); i++ {
		now := time.Now()
		creationDate, err := time.Parse(time.RFC3339Nano, *resp.Images[i].CreationDate)
		if err != nil {
			panic(err)
		}

		duration := now.Sub(creationDate)

		if duration.Hours() > hours {
			filteredAmis = append(filteredAmis, resp.Images[i])
		}
	}
	c.Ui.Output("==> Found " + strconv.Itoa(len(filteredAmis)) + " total AMI(s) for deletion.")

	if len(filteredAmis) == 0 {
		c.Ui.Error("No AMIs to delete.")
		return 0
	}

	// Get a list of every single snapshot in our account
	// (I wasn't able to find a better way to filter these, but suggestions welcome!)
	respDscrSnapshots, err := svc.DescribeSnapshots(&ec2.DescribeSnapshotsInput{
		OwnerIDs: []*string{&awsAccountId},
	})
	if err != nil {
		panic(err)
	}
	c.Ui.Output("==> Found " + strconv.Itoa(len(respDscrSnapshots.Snapshots)) + " total snapshots in this account.")

	// Compute whether we should delete fewer AMIs to adhere to our --require-at-least requirement
	var numTotalAmis = len(resp.Images)
	var numFilteredAmis = len(filteredAmis)
	var numAmisToRemainAfterDelete = numTotalAmis - numFilteredAmis
	var numAmisToRemoveFromFiltered = math.Max(0.0, float64(c.RequireAtLeast - numAmisToRemainAfterDelete))

	if numAmisToRemoveFromFiltered > 0.0 {
		c.Ui.Output("==> Only deleting " + strconv.Itoa(len(filteredAmis) - int(numAmisToRemoveFromFiltered)) + " total AMIs to honor '--require-at-least=" + strconv.Itoa(c.RequireAtLeast) + "'.")
	}

	// Begin deleting AMIs...
	for i := 0; i < len(filteredAmis) - int(numAmisToRemoveFromFiltered); i++ {
		// Step 1: De-register the AMI
		c.Ui.Output(*filteredAmis[i].ImageID + ": De-registering AMI named \"" + *filteredAmis[i].Name + "\"...")
		_, err := svc.DeregisterImage(&ec2.DeregisterImageInput{
			DryRun: &c.DryRun,
			ImageID: filteredAmis[i].ImageID,
		})
		if err != nil {
			if ! strings.Contains(err.Error(), "DryRunOperation") {
				panic(err)
			}
		}

		// Step 2: Delete the corresponding AMI snapshot
		// Look at the "description" for each Snapshot to see if it contains our AMI id
		var snapshotIds []string
		for _, snapshot := range respDscrSnapshots.Snapshots {
			if strings.Contains(*snapshot.Description, *filteredAmis[i].ImageID) {
				snapshotIds = append(snapshotIds, *snapshot.SnapshotID)
			}
		}

		// Delete all snapshots that were found
		c.Ui.Output(*filteredAmis[i].ImageID + ": Found " + strconv.Itoa(len(snapshotIds)) + " snapshot(s) to delete")
		for _, snapshotId := range snapshotIds {
			c.Ui.Output(*filteredAmis[i].ImageID + ": Deleting snapshot " + snapshotId + "...")
			svc.DeleteSnapshot(&ec2.DeleteSnapshotInput{
				DryRun: &c.DryRun,
				SnapshotID: &snapshotId,
			})
		}

		c.Ui.Output(*filteredAmis[i].ImageID + ": Done!")
		c.Ui.Output("")
	}

	if c.DryRun {
		c.Ui.Info("==> DRY RUN. Had this not been a dry run, " + strconv.Itoa(len(filteredAmis)) + " AMI's and their corresponding snapshots would have been deleted.")
	} else {
		c.Ui.Info("==> Success! Deleted " + strconv.Itoa(len(filteredAmis) - int(numAmisToRemoveFromFiltered)) + " AMI's and their corresponding snapshots.")
	}
	return 0
}

<|endoftext|>


<commit_before>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package changelist

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"
)

// CLID is a unique ID of a CL used internally in CV.
//
// It's just 8 bytes long and is thus much shorter than ExternalID,
// which reduces CPU & RAM & storage costs of CL graphs for multi-CL Runs.
type CLID int64

// ExternalID is a unique CL ID deterministically constructed based on CL data.
//
// Currently, only Gerrit is supported.
type ExternalID string

// GobID makes an ExternalID for a Gerrit CL.
//
// Host is typically "something-review.googlesource.com".
// Change is a number, e.g. 2515619 for
// https://chromium-review.googlesource.com/c/infra/luci/luci-go/+/2515619
func GobID(host string, change int64) (ExternalID, error) {
	if strings.ContainsRune(host, '/') {
		return "", errors.Reason("invalid host %q: must not contain /", host).Err()
	}
	return ExternalID(fmt.Sprintf("gerrit/%s/%d", host, change)), nil
}

// ParseGobID returns Gerrit host and change if this is a GobID.
func (e ExternalID) ParseGobID() (host string, change int64, err error) {
	parts := strings.Split(string(e), "/")
	if len(parts) != 3 || parts[0] != "gerrit" {
		err = errors.Reason("%q is not a valid GobID", e).Err()
		return
	}
	host = parts[1]
	change, err = strconv.ParseInt(parts[2], 10, 63)
	if err != nil {
		err = errors.Annotate(err, "%q is not a valid GobID", e).Err()
	}
	return
}

// CL is a CL entity in Datastore.
type CL struct {
	_kind  string                `gae:"$kind,CL"`
	_extra datastore.PropertyMap `gae:"-,extra"`

	// ID is auto-generated by Datastore.
	ID CLID `gae:"$id"` // int64
	// ExternalID must not be modified once entity is created.
	ExternalID ExternalID `gae:",noindex"` // string. Indexed in CLMap entities.

	// EVersion is entity version. Every update should increment it by 1.
	// See Update() function.
	EVersion int `gae:",noindex"`

	// Snapshot is latest known state of a CL.
	// It may and often is behind the source of truth -- the code reveview site
	// (e.g. Gerrit).
	Snapshot *Snapshot

	// ApplicableConfig keeps track of configs applicable to the CL.
	ApplicableConfig *ApplicableConfig

	// UpdateTime is exact time of when this entity was last updated.
	//
	// It's not indexed to avoid hot areas in the index.
	UpdateTime time.Time `gae:",noindex"`

	// TODO(tandrii): implement deletion of the oldest entities via additional
	// indexed field based on UpdateTime but with entropy in the lowest bits to
	// avoid hotspots.
}

// clMap is CLMap entity in Datastore which ensures strict 1:1 mapping
// between internal and external IDs.
type clMap struct {
	_kind string `gae:"$kind,CLMap"`

	// ExternalID as entity ID ensures uniqueness.
	ExternalID ExternalID `gae:"$id"` // string
	// InternalID is auto-generated by Datastore for CL entity.
	InternalID CLID `gae:",noindex"` // int64. Indexed in CL entities.
}

// Get reads a CL from datastore.
//
// Returns datastore.ErrNoSuchEntity if it doesn't exist.
func (eid ExternalID) Get(ctx context.Context) (*CL, error) {
	m := clMap{ExternalID: eid}
	switch err := datastore.Get(ctx, &m); {
	case err == datastore.ErrNoSuchEntity:
		return nil, err
	case err != nil:
		return nil, errors.Annotate(err, "failed to get CLMap").Tag(transient.Tag).Err()
	}
	return getExisting(ctx, m.InternalID, eid)
}

// GetOrInsert reads a CL from datastore, creating a new one if not exists yet.
//
// populate is called within a transaction to populate fields of a new entity.
// It should be a fast function.
//
// Warning:
//  * populate may be called several times since transaction can be retried.
//  * cl.ExternalID and cl.ID must not be changed by populate.
func (eid ExternalID) GetOrInsert(ctx context.Context, populate func(cl *CL)) (*CL, error) {
	// Fast path without transaction.
	if cl, err := eid.Get(ctx); err != datastore.ErrNoSuchEntity {
		return cl, err
	}
	var cl *CL
	m := clMap{ExternalID: eid}
	err := datastore.RunInTransaction(ctx, func(ctx context.Context) (err error) {
		cl = nil
		switch err = datastore.Get(ctx, &m); {
		case err == nil:
			// Has just been created by someone else.
			return nil
		case err != datastore.ErrNoSuchEntity:
			return err
		}
		cl, err = insert(ctx, eid, populate)
		return
	}, nil)

	switch {
	case err != nil:
		return nil, errors.Annotate(err, "failed to getOrInsert a CL").Tag(transient.Tag).Err()
	case cl == nil:
		return getExisting(ctx, m.InternalID, eid)
	}
	return cl, nil
}

// Delete deletes CL and its CLMap entities trasactionally.
//
// Thus, Delete and insertion (part of ExternalID.getOrInsert)
// are atomic with respect to one another.
//
// However, ExternalID.get and fast path of ExternalID.getOrInsert if called
// concurrently with Delete may return temporary error, but on retry they would
// return ErrNoSuchEntity.
func Delete(ctx context.Context, id CLID) error {
	cl := CL{ID: id}
	switch err := datastore.Get(ctx, &cl); {
	case err == datastore.ErrNoSuchEntity:
		return nil // Nothing to do.
	case err != nil:
		return errors.Annotate(err, "failed to get a CL").Tag(transient.Tag).Err()
	}

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		m := clMap{ExternalID: cl.ExternalID}
		return datastore.Delete(ctx, &cl, &m)
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to delete a CL").Tag(transient.Tag).Err()
	}
	return nil
}

// Update updates CL entity with Snapshot and ApplicableConfig.
//
// Either ExternalID or a known CLID must be provided.
// Either new Snapshot or ApplicableConfig must be provided.
//
// If CLID is not known and CL for provided ExternalID doesn't exist,
// then a new CL is created with the given Snapshot & ApplicableConfig.
//
// Otherwise, an existing CL entity will be updated iff either:
//  * if snapshot is given and it is more recent than what is already stored,
//    as measured by ExternalUpdateTime.
//  * same but for ApplicableConfig and ApplicableConfig.UpdateTime,
//    respectively.
//
// TODO(tandrii): emit notification events.
func Update(ctx context.Context, eid ExternalID, knownCLID CLID, snapshot *Snapshot, acfg *ApplicableConfig) error {
	if eid == "" && knownCLID == 0 {
		panic("either ExternalID or known CLID must be provided")
	}
	if snapshot == nil && acfg == nil {
		panic("either new snapshot or new ApplicableConfig must be provided")
	}

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		if knownCLID == 0 {
			m := clMap{ExternalID: eid}
			switch err := datastore.Get(ctx, &m); {
			case err == datastore.ErrNoSuchEntity:
				// Insert new entity.
				_, err = insert(ctx, eid, func(cl *CL) {
					cl.Snapshot = snapshot
					cl.ApplicableConfig = acfg
				})
				return err
			case err != nil:
				return errors.Annotate(err, "failed to get CLMap entity").Tag(transient.Tag).Err()
			}
			knownCLID = m.InternalID
		}
		cl, err := getExisting(ctx, knownCLID, eid)
		if err != nil {
			return err
		}
		// Update exsting entity.
		return update(ctx, cl, func(cl *CL) (changed bool) {
			if snapshot != nil && !cl.Snapshot.IsUpToDate(snapshot.GetLuciProject(), snapshot.GetExternalUpdateTime().AsTime()) {
				cl.Snapshot = snapshot
				changed = true
			}
			if acfg != nil && !cl.ApplicableConfig.IsUpToDate(acfg.GetUpdateTime().AsTime()) {
				cl.ApplicableConfig = acfg
				changed = true
			}
			return
		})
	}, nil)
	return errors.Annotate(err, "failed to update CL").Tag(transient.Tag).Err()
}

func getExisting(ctx context.Context, clid CLID, eid ExternalID) (*CL, error) {
	cl := &CL{ID: clid}
	switch err := datastore.Get(ctx, cl); {
	case err == datastore.ErrNoSuchEntity:
		// This should not happen in practice except in the case of a very old CL
		// which is being deleted due to retention policy. Log error but return it
		// as transient as it's expected that CLMap entity would be removed soon,
		// and so a retry would be produce proper datastore.ErrNoSuchEntity error.
		msg := fmt.Sprintf("unexpectedly failed to get CL#%d given existing CLMap%q", clid, eid)
		logging.Errorf(ctx, msg)
		return nil, errors.Reason(msg).Tag(transient.Tag).Err()
	case err != nil:
		return nil, errors.Annotate(err, "failed to get CL").Tag(transient.Tag).Err()
	}
	return cl, nil
}

// insert creates new CL entity for given external ID.
//
// Must be called after verifying that such CLMap record doesn't exist.
func insert(ctx context.Context, eid ExternalID, populate func(*CL)) (*CL, error) {
	if datastore.CurrentTransaction(ctx) == nil {
		panic("must be called in transaction context")
	}
	// Create new CL and CLMap entry atomically.
	cl := &CL{
		ID:         0, // autogenerate by Datastore
		ExternalID: eid,
		EVersion:   1,
	}
	populate(cl)
	if cl.ID != 0 || cl.ExternalID != eid || cl.EVersion != 1 {
		panic(errors.New("populate changed ID or ExternalID or EVersion, but must not do this."))
	}
	cl.UpdateTime = datastore.RoundTime(clock.Now(ctx).UTC())

	if err := datastore.Put(ctx, cl); err != nil {
		return nil, errors.Annotate(err, "failed to save CL entity").Tag(transient.Tag).Err()
	}
	if err := datastore.Put(ctx, &clMap{ExternalID: eid, InternalID: cl.ID}); err != nil {
		return nil, errors.Annotate(err, "failed to save CLMap entity").Tag(transient.Tag).Err()
	}
	return cl, nil
}

func update(ctx context.Context, justRead *CL, mut func(*CL) (update bool)) error {
	if datastore.CurrentTransaction(ctx) == nil {
		panic("must be called in transaction context")
	}

	before := *justRead // shallow copy, avoiding cloning Snapshot.
	if !mut(justRead) {
		return nil
	}
	justRead.EVersion = before.EVersion + 1
	justRead.UpdateTime = clock.Now(ctx).UTC()
	if err := datastore.Put(ctx, justRead); err != nil {
		return errors.Annotate(err, "failed to put CL entity").Tag(transient.Tag).Err()
	}
	return nil
}
<commit_msg>cv: explain why datastore.RoundTime is called before datastore.Put.<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package changelist

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"
	"go.chromium.org/luci/gae/service/datastore"
)

// CLID is a unique ID of a CL used internally in CV.
//
// It's just 8 bytes long and is thus much shorter than ExternalID,
// which reduces CPU & RAM & storage costs of CL graphs for multi-CL Runs.
type CLID int64

// ExternalID is a unique CL ID deterministically constructed based on CL data.
//
// Currently, only Gerrit is supported.
type ExternalID string

// GobID makes an ExternalID for a Gerrit CL.
//
// Host is typically "something-review.googlesource.com".
// Change is a number, e.g. 2515619 for
// https://chromium-review.googlesource.com/c/infra/luci/luci-go/+/2515619
func GobID(host string, change int64) (ExternalID, error) {
	if strings.ContainsRune(host, '/') {
		return "", errors.Reason("invalid host %q: must not contain /", host).Err()
	}
	return ExternalID(fmt.Sprintf("gerrit/%s/%d", host, change)), nil
}

// ParseGobID returns Gerrit host and change if this is a GobID.
func (e ExternalID) ParseGobID() (host string, change int64, err error) {
	parts := strings.Split(string(e), "/")
	if len(parts) != 3 || parts[0] != "gerrit" {
		err = errors.Reason("%q is not a valid GobID", e).Err()
		return
	}
	host = parts[1]
	change, err = strconv.ParseInt(parts[2], 10, 63)
	if err != nil {
		err = errors.Annotate(err, "%q is not a valid GobID", e).Err()
	}
	return
}

// CL is a CL entity in Datastore.
type CL struct {
	_kind  string                `gae:"$kind,CL"`
	_extra datastore.PropertyMap `gae:"-,extra"`

	// ID is auto-generated by Datastore.
	ID CLID `gae:"$id"` // int64
	// ExternalID must not be modified once entity is created.
	ExternalID ExternalID `gae:",noindex"` // string. Indexed in CLMap entities.

	// EVersion is entity version. Every update should increment it by 1.
	// See Update() function.
	EVersion int `gae:",noindex"`

	// Snapshot is latest known state of a CL.
	// It may and often is behind the source of truth -- the code reveview site
	// (e.g. Gerrit).
	Snapshot *Snapshot

	// ApplicableConfig keeps track of configs applicable to the CL.
	ApplicableConfig *ApplicableConfig

	// UpdateTime is exact time of when this entity was last updated.
	//
	// It's not indexed to avoid hot areas in the index.
	UpdateTime time.Time `gae:",noindex"`

	// TODO(tandrii): implement deletion of the oldest entities via additional
	// indexed field based on UpdateTime but with entropy in the lowest bits to
	// avoid hotspots.
}

// clMap is CLMap entity in Datastore which ensures strict 1:1 mapping
// between internal and external IDs.
type clMap struct {
	_kind string `gae:"$kind,CLMap"`

	// ExternalID as entity ID ensures uniqueness.
	ExternalID ExternalID `gae:"$id"` // string
	// InternalID is auto-generated by Datastore for CL entity.
	InternalID CLID `gae:",noindex"` // int64. Indexed in CL entities.
}

// Get reads a CL from datastore.
//
// Returns datastore.ErrNoSuchEntity if it doesn't exist.
func (eid ExternalID) Get(ctx context.Context) (*CL, error) {
	m := clMap{ExternalID: eid}
	switch err := datastore.Get(ctx, &m); {
	case err == datastore.ErrNoSuchEntity:
		return nil, err
	case err != nil:
		return nil, errors.Annotate(err, "failed to get CLMap").Tag(transient.Tag).Err()
	}
	return getExisting(ctx, m.InternalID, eid)
}

// GetOrInsert reads a CL from datastore, creating a new one if not exists yet.
//
// populate is called within a transaction to populate fields of a new entity.
// It should be a fast function.
//
// Warning:
//  * populate may be called several times since transaction can be retried.
//  * cl.ExternalID and cl.ID must not be changed by populate.
func (eid ExternalID) GetOrInsert(ctx context.Context, populate func(cl *CL)) (*CL, error) {
	// Fast path without transaction.
	if cl, err := eid.Get(ctx); err != datastore.ErrNoSuchEntity {
		return cl, err
	}
	var cl *CL
	m := clMap{ExternalID: eid}
	err := datastore.RunInTransaction(ctx, func(ctx context.Context) (err error) {
		cl = nil
		switch err = datastore.Get(ctx, &m); {
		case err == nil:
			// Has just been created by someone else.
			return nil
		case err != datastore.ErrNoSuchEntity:
			return err
		}
		cl, err = insert(ctx, eid, populate)
		return
	}, nil)

	switch {
	case err != nil:
		return nil, errors.Annotate(err, "failed to getOrInsert a CL").Tag(transient.Tag).Err()
	case cl == nil:
		return getExisting(ctx, m.InternalID, eid)
	}
	return cl, nil
}

// Delete deletes CL and its CLMap entities trasactionally.
//
// Thus, Delete and insertion (part of ExternalID.getOrInsert)
// are atomic with respect to one another.
//
// However, ExternalID.get and fast path of ExternalID.getOrInsert if called
// concurrently with Delete may return temporary error, but on retry they would
// return ErrNoSuchEntity.
func Delete(ctx context.Context, id CLID) error {
	cl := CL{ID: id}
	switch err := datastore.Get(ctx, &cl); {
	case err == datastore.ErrNoSuchEntity:
		return nil // Nothing to do.
	case err != nil:
		return errors.Annotate(err, "failed to get a CL").Tag(transient.Tag).Err()
	}

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		m := clMap{ExternalID: cl.ExternalID}
		return datastore.Delete(ctx, &cl, &m)
	}, nil)
	if err != nil {
		return errors.Annotate(err, "failed to delete a CL").Tag(transient.Tag).Err()
	}
	return nil
}

// Update updates CL entity with Snapshot and ApplicableConfig.
//
// Either ExternalID or a known CLID must be provided.
// Either new Snapshot or ApplicableConfig must be provided.
//
// If CLID is not known and CL for provided ExternalID doesn't exist,
// then a new CL is created with the given Snapshot & ApplicableConfig.
//
// Otherwise, an existing CL entity will be updated iff either:
//  * if snapshot is given and it is more recent than what is already stored,
//    as measured by ExternalUpdateTime.
//  * same but for ApplicableConfig and ApplicableConfig.UpdateTime,
//    respectively.
//
// TODO(tandrii): emit notification events.
func Update(ctx context.Context, eid ExternalID, knownCLID CLID, snapshot *Snapshot, acfg *ApplicableConfig) error {
	if eid == "" && knownCLID == 0 {
		panic("either ExternalID or known CLID must be provided")
	}
	if snapshot == nil && acfg == nil {
		panic("either new snapshot or new ApplicableConfig must be provided")
	}

	err := datastore.RunInTransaction(ctx, func(ctx context.Context) error {
		if knownCLID == 0 {
			m := clMap{ExternalID: eid}
			switch err := datastore.Get(ctx, &m); {
			case err == datastore.ErrNoSuchEntity:
				// Insert new entity.
				_, err = insert(ctx, eid, func(cl *CL) {
					cl.Snapshot = snapshot
					cl.ApplicableConfig = acfg
				})
				return err
			case err != nil:
				return errors.Annotate(err, "failed to get CLMap entity").Tag(transient.Tag).Err()
			}
			knownCLID = m.InternalID
		}
		cl, err := getExisting(ctx, knownCLID, eid)
		if err != nil {
			return err
		}
		// Update exsting entity.
		return update(ctx, cl, func(cl *CL) (changed bool) {
			if snapshot != nil && !cl.Snapshot.IsUpToDate(snapshot.GetLuciProject(), snapshot.GetExternalUpdateTime().AsTime()) {
				cl.Snapshot = snapshot
				changed = true
			}
			if acfg != nil && !cl.ApplicableConfig.IsUpToDate(acfg.GetUpdateTime().AsTime()) {
				cl.ApplicableConfig = acfg
				changed = true
			}
			return
		})
	}, nil)
	return errors.Annotate(err, "failed to update CL").Tag(transient.Tag).Err()
}

func getExisting(ctx context.Context, clid CLID, eid ExternalID) (*CL, error) {
	cl := &CL{ID: clid}
	switch err := datastore.Get(ctx, cl); {
	case err == datastore.ErrNoSuchEntity:
		// This should not happen in practice except in the case of a very old CL
		// which is being deleted due to retention policy. Log error but return it
		// as transient as it's expected that CLMap entity would be removed soon,
		// and so a retry would be produce proper datastore.ErrNoSuchEntity error.
		msg := fmt.Sprintf("unexpectedly failed to get CL#%d given existing CLMap%q", clid, eid)
		logging.Errorf(ctx, msg)
		return nil, errors.Reason(msg).Tag(transient.Tag).Err()
	case err != nil:
		return nil, errors.Annotate(err, "failed to get CL").Tag(transient.Tag).Err()
	}
	return cl, nil
}

// insert creates new CL entity for given external ID.
//
// Must be called after verifying that such CLMap record doesn't exist.
func insert(ctx context.Context, eid ExternalID, populate func(*CL)) (*CL, error) {
	if datastore.CurrentTransaction(ctx) == nil {
		panic("must be called in transaction context")
	}
	// Create new CL and CLMap entry atomically.
	cl := &CL{
		ID:         0, // autogenerate by Datastore
		ExternalID: eid,
		EVersion:   1,
	}
	populate(cl)
	if cl.ID != 0 || cl.ExternalID != eid || cl.EVersion != 1 {
		panic(errors.New("populate changed ID or ExternalID or EVersion, but must not do this."))
	}
	// datastore.Put will do RoundTime on its own, but without affecting our `cl`
	// object. Since `cl` object is passed outside, do rounding here s.t. it has
	// exact same data as would have been read from datastore right after the Put.
	cl.UpdateTime = datastore.RoundTime(clock.Now(ctx).UTC())

	if err := datastore.Put(ctx, cl); err != nil {
		return nil, errors.Annotate(err, "failed to save CL entity").Tag(transient.Tag).Err()
	}
	if err := datastore.Put(ctx, &clMap{ExternalID: eid, InternalID: cl.ID}); err != nil {
		return nil, errors.Annotate(err, "failed to save CLMap entity").Tag(transient.Tag).Err()
	}
	return cl, nil
}

func update(ctx context.Context, justRead *CL, mut func(*CL) (update bool)) error {
	if datastore.CurrentTransaction(ctx) == nil {
		panic("must be called in transaction context")
	}

	before := *justRead // shallow copy, avoiding cloning Snapshot.
	if !mut(justRead) {
		return nil
	}
	justRead.EVersion = before.EVersion + 1
	justRead.UpdateTime = clock.Now(ctx).UTC()
	if err := datastore.Put(ctx, justRead); err != nil {
		return errors.Annotate(err, "failed to put CL entity").Tag(transient.Tag).Err()
	}
	return nil
}
<|endoftext|>


<commit_before>package varmor

import (
	"math/rand"
	"testing"

	"github.com/stretchr/testify/assert"
)

func preserveString(t *testing.T, s string) {
	b, err := Unwrap(Wrap([]byte(s)))
	assert.NoError(t, err)
	assert.Equal(t, s, string(b))
}

func preserveBytes(t *testing.T, b []byte) {
	wrapped, err := Unwrap(Wrap(b))
	assert.NoError(t, err)
	assert.EqualValues(t, b, wrapped)
}

func TestPreservation(t *testing.T) {
	preserveString(t, "")
	preserveString(t, "test")

	rnd := rand.New(rand.NewSource(0))
	rbytes := make([]byte, 100000)
	n, err := rnd.Read(rbytes)
	assert.NoError(t, err)
	assert.Equal(t, 100000, n)
	preserveBytes(t, rbytes)
}

func TestTruncated(t *testing.T) {
	b, err := Unwrap("")
	assert.Error(t, err)
	assert.Nil(t, b)
}

func TestWrongVersion(t *testing.T) {
	b, err := Unwrap("saltybox999999:...")
	assert.Error(t, err)
	assert.Equal(t, "input claims to be saltybox, but not a version we support", err.Error())
	assert.Nil(t, b)
}

func TestAllByteValues(t *testing.T) {
	allBytes := make([]byte, 256)
	for i := 0; i <= 255; i++ {
		allBytes[i] = byte(i)
	}

	preserveBytes(t, allBytes)

	wrapped := Wrap(allBytes)
	assert.Equal(t,
		"saltybox1:AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8gISIjJCUmJygpKissLS4vMDEyMzQ1Njc4OTo7PD0-P0BBQkNERUZHSElKS0xNTk9QUVJTVFVWV1hZWltcXV5fYGFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6e3x9fn-AgYKDhIWGh4iJiouMjY6PkJGSk5SVlpeYmZqbnJ2en6ChoqOkpaanqKmqq6ytrq-wsbKztLW2t7i5uru8vb6_wMHCw8TFxsfIycrLzM3Oz9DR0tPU1dbX2Nna29zd3t_g4eLj5OXm5-jp6uvs7e7v8PHy8_T19vf4-fr7_P3-_w",
		wrapped)
}
<commit_msg>Add test case. (#48)<commit_after>package varmor

import (
	"math/rand"
	"testing"

	"github.com/stretchr/testify/assert"
)

func preserveString(t *testing.T, s string) {
	b, err := Unwrap(Wrap([]byte(s)))
	assert.NoError(t, err)
	assert.Equal(t, s, string(b))
}

func preserveBytes(t *testing.T, b []byte) {
	wrapped, err := Unwrap(Wrap(b))
	assert.NoError(t, err)
	assert.EqualValues(t, b, wrapped)
}

func TestPreservation(t *testing.T) {
	preserveString(t, "")
	preserveString(t, "test")

	rnd := rand.New(rand.NewSource(0))
	rbytes := make([]byte, 100000)
	n, err := rnd.Read(rbytes)
	assert.NoError(t, err)
	assert.Equal(t, 100000, n)
	preserveBytes(t, rbytes)
}

func TestTruncated(t *testing.T) {
	b, err := Unwrap("")
	assert.Error(t, err)
	assert.Nil(t, b)
}

func TestWrongVersion(t *testing.T) {
	b, err := Unwrap("saltybox999999:...")
	assert.Error(t, err)
	assert.Equal(t, "input claims to be saltybox, but not a version we support", err.Error())
	assert.Nil(t, b)
}

func TestNotSaltybox(t *testing.T) {
	b, err := Unwrap("something not looking like saltybox data")
	assert.Error(t, err)
	assert.Equal(t, "input unrecognized as saltybox data", err.Error())
	assert.Nil(t, b)
}

func TestAllByteValues(t *testing.T) {
	allBytes := make([]byte, 256)
	for i := 0; i <= 255; i++ {
		allBytes[i] = byte(i)
	}

	preserveBytes(t, allBytes)

	wrapped := Wrap(allBytes)
	assert.Equal(t,
		"saltybox1:AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8gISIjJCUmJygpKissLS4vMDEyMzQ1Njc4OTo7PD0-P0BBQkNERUZHSElKS0xNTk9QUVJTVFVWV1hZWltcXV5fYGFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6e3x9fn-AgYKDhIWGh4iJiouMjY6PkJGSk5SVlpeYmZqbnJ2en6ChoqOkpaanqKmqq6ytrq-wsbKztLW2t7i5uru8vb6_wMHCw8TFxsfIycrLzM3Oz9DR0tPU1dbX2Nna29zd3t_g4eLj5OXm5-jp6uvs7e7v8PHy8_T19vf4-fr7_P3-_w",
		wrapped)
}
<|endoftext|>


<commit_before>package topgun_test

import (
	"bytes"
	"os"
	"regexp"
	"time"

	_ "github.com/lib/pq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("[#129726011] Worker landing", func() {
	Context("with two workers available", func() {
		BeforeEach(func() {
			Deploy("deployments/two-forwarded-workers.yml")
		})

		Describe("restarting the worker", func() {
			var restartingWorkerName string
			var restartSession *gexec.Session

			JustBeforeEach(func() {
				restartSession = spawnBosh("restart", "worker/0")
				restartingWorkerName = waitForLandingOrLandedWorker()
			})

			AfterEach(func() {
				<-restartSession.Exited
			})

			Context("while in landing or landed state", func() {
				// technically this is timing-dependent but it doesn't seem worth the
				// time cost of explicit tests for both

				It("is not used for new workloads", func() {
					for i := 0; i < 10; i++ {
						fly("execute", "-c", "tasks/tiny.yml")
						usedWorkers := workersWithContainers()
						Expect(usedWorkers).To(HaveLen(1))
						Expect(usedWorkers).ToNot(ContainElement(restartingWorkerName))
					}
				})

				It("can be pruned", func() {
					fly("prune-worker", "-w", restartingWorkerName)
					waitForWorkersToBeRunning()
				})
			})
		})
	})

	describeRestartingTheWorker := func() {
		Describe("restarting the worker", func() {
			var restartSession *gexec.Session

			JustBeforeEach(func() {
				restartSession = spawnBosh("restart", "worker/0")
				_ = waitForLandingOrLandedWorker()
			})

			AfterEach(func() {
				<-restartSession.Exited
			})

			Context("with volumes and containers present", func() {
				var preservedContainerID string

				BeforeEach(func() {
					By("setting pipeline that creates volumes for image")
					fly("set-pipeline", "-n", "-c", "pipelines/get-task.yml", "-p", "topgun")

					By("unpausing the pipeline")
					fly("unpause-pipeline", "-p", "topgun")

					By("triggering a job")
					buildSession := spawnFly("trigger-job", "-w", "-j", "topgun/simple-job")
					Eventually(buildSession).Should(gbytes.Say("Pulling .*busybox.*"))
					<-buildSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("getting identifier for check container")
					hijackSession := spawnFly("hijack", "-c", "topgun/tick-tock", "--", "hostname")
					<-hijackSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					preservedContainerID = string(hijackSession.Out.Contents())
				})

				It("keeps volumes and containers after restart", func() {
					By("completing the restart")
					<-restartSession.Exited
					Expect(restartSession.ExitCode()).To(Equal(0))

					By("retaining cached image resource in second job build")
					buildSession := spawnFly("trigger-job", "-w", "-j", "topgun/simple-job")
					<-buildSession.Exited
					Expect(buildSession).NotTo(gbytes.Say("Pulling .*busybox.*"))
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("retaining check containers")
					hijackSession := spawnFly("hijack", "-c", "topgun/tick-tock", "--", "hostname")
					<-hijackSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					currentContainerID := string(hijackSession.Out.Contents())
					Expect(currentContainerID).To(Equal(preservedContainerID))
				})
			})

			Context("with an interruptible build in-flight", func() {
				var buildSession *gexec.Session

				BeforeEach(func() {
					By("setting pipeline that has an infinite but interruptible job")
					fly("set-pipeline", "-n", "-c", "pipelines/interruptible.yml", "-p", "topgun")

					By("unpausing the pipeline")
					fly("unpause-pipeline", "-p", "topgun")

					By("triggering a job")
					buildSession = spawnFly("trigger-job", "-w", "-j", "topgun/interruptible-job")
					Eventually(buildSession).Should(gbytes.Say("waiting forever"))
				})

				It("does not wait for the build", func() {
					By("completing the restart without the drain timeout kicking in")
					Eventually(restartSession, 5*time.Minute).Should(gexec.Exit(0))
				})
			})

			Context("with uninterruptible build in-flight", func() {
				var buildSession *gexec.Session
				var buildID string

				BeforeEach(func() {
					buildSession = spawnFly("execute", "-c", "tasks/wait.yml")
					Eventually(buildSession).Should(gbytes.Say("executing build"))

					buildRegex := regexp.MustCompile(`executing build (\d+)`)
					matches := buildRegex.FindSubmatch(buildSession.Out.Contents())
					buildID = string(matches[1])

					Eventually(buildSession).Should(gbytes.Say("waiting for /tmp/stop-waiting"))
				})

				AfterEach(func() {
					buildSession.Signal(os.Interrupt)
					<-buildSession.Exited
				})

				It("waits for the build", func() {
					Eventually(restartSession).Should(gbytes.Say(`Updating (instance|job)`))
					Consistently(restartSession, 5*time.Minute).ShouldNot(gexec.Exit())
				})

				It("finishes restarting once the build is done", func() {
					By("hijacking the build to tell it to finish")
					<-flyHijackTask(
						"-b", buildID,
						"-s", "one-off",
						"touch", "/tmp/stop-waiting",
					).Exited

					By("waiting for the build to exit")
					Eventually(buildSession).Should(gbytes.Say("done"))
					<-buildSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("successfully restarting")
					<-restartSession.Exited
					Expect(restartSession.ExitCode()).To(Equal(0))
				})
			})
		})
	}

	Context("with one worker", func() {
		BeforeEach(func() {
			Deploy("deployments/one-forwarded-worker.yml")
		})

		describeRestartingTheWorker()
	})

	Context("with a single team worker", func() {
		BeforeEach(func() {
			Deploy("deployments/team-worker.yml")

			setTeam := spawnFlyInteractive(bytes.NewBufferString("y\n"), "set-team", "-n", "team-a", "--no-really-i-dont-want-any-auth")
			<-setTeam.Exited
			Expect(setTeam.ExitCode()).To(Equal(0))

			fly("login", "-c", atcExternalURL, "-n", "team-a")

			// wait for the team's worker to arrive now that team exists
			waitForRunningWorker()
		})

		describeRestartingTheWorker()
	})
})
<commit_msg>wait for the worker to be running<commit_after>package topgun_test

import (
	"bytes"
	"os"
	"regexp"
	"time"

	_ "github.com/lib/pq"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gbytes"
	"github.com/onsi/gomega/gexec"
)

var _ = Describe("[#129726011] Worker landing", func() {
	Context("with two workers available", func() {
		BeforeEach(func() {
			Deploy("deployments/two-forwarded-workers.yml")
		})

		Describe("restarting the worker", func() {
			var restartingWorkerName string
			var restartSession *gexec.Session

			JustBeforeEach(func() {
				restartSession = spawnBosh("restart", "worker/0")
				restartingWorkerName = waitForLandingOrLandedWorker()
			})

			AfterEach(func() {
				<-restartSession.Exited
			})

			Context("while in landing or landed state", func() {
				// technically this is timing-dependent but it doesn't seem worth the
				// time cost of explicit tests for both

				It("is not used for new workloads", func() {
					for i := 0; i < 10; i++ {
						fly("execute", "-c", "tasks/tiny.yml")
						usedWorkers := workersWithContainers()
						Expect(usedWorkers).To(HaveLen(1))
						Expect(usedWorkers).ToNot(ContainElement(restartingWorkerName))
					}
				})

				It("can be pruned", func() {
					fly("prune-worker", "-w", restartingWorkerName)
					waitForWorkersToBeRunning()
				})
			})
		})
	})

	describeRestartingTheWorker := func() {
		Describe("restarting the worker", func() {
			var restartSession *gexec.Session

			JustBeforeEach(func() {
				restartSession = spawnBosh("restart", "worker/0")
				_ = waitForLandingOrLandedWorker()
			})

			AfterEach(func() {
				<-restartSession.Exited
			})

			Context("with volumes and containers present", func() {
				var preservedContainerID string

				BeforeEach(func() {
					By("setting pipeline that creates volumes for image")
					fly("set-pipeline", "-n", "-c", "pipelines/get-task.yml", "-p", "topgun")

					By("unpausing the pipeline")
					fly("unpause-pipeline", "-p", "topgun")

					By("triggering a job")
					buildSession := spawnFly("trigger-job", "-w", "-j", "topgun/simple-job")
					Eventually(buildSession).Should(gbytes.Say("Pulling .*busybox.*"))
					<-buildSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("getting identifier for check container")
					hijackSession := spawnFly("hijack", "-c", "topgun/tick-tock", "--", "hostname")
					<-hijackSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					preservedContainerID = string(hijackSession.Out.Contents())
				})

				It("keeps volumes and containers after restart", func() {
					By("completing the restart")
					<-restartSession.Exited
					Expect(restartSession.ExitCode()).To(Equal(0))

					By("retaining cached image resource in second job build")
					buildSession := spawnFly("trigger-job", "-w", "-j", "topgun/simple-job")
					<-buildSession.Exited
					Expect(buildSession).NotTo(gbytes.Say("Pulling .*busybox.*"))
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("retaining check containers")
					hijackSession := spawnFly("hijack", "-c", "topgun/tick-tock", "--", "hostname")
					<-hijackSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					currentContainerID := string(hijackSession.Out.Contents())
					Expect(currentContainerID).To(Equal(preservedContainerID))
				})
			})

			Context("with an interruptible build in-flight", func() {
				var buildSession *gexec.Session

				BeforeEach(func() {
					By("setting pipeline that has an infinite but interruptible job")
					fly("set-pipeline", "-n", "-c", "pipelines/interruptible.yml", "-p", "topgun")

					By("unpausing the pipeline")
					fly("unpause-pipeline", "-p", "topgun")

					By("triggering a job")
					buildSession = spawnFly("trigger-job", "-w", "-j", "topgun/interruptible-job")
					Eventually(buildSession).Should(gbytes.Say("waiting forever"))
				})

				It("does not wait for the build", func() {
					By("completing the restart without the drain timeout kicking in")
					Eventually(restartSession, 5*time.Minute).Should(gexec.Exit(0))
				})
			})

			Context("with uninterruptible build in-flight", func() {
				var buildSession *gexec.Session
				var buildID string

				BeforeEach(func() {
					buildSession = spawnFly("execute", "-c", "tasks/wait.yml")
					Eventually(buildSession).Should(gbytes.Say("executing build"))

					buildRegex := regexp.MustCompile(`executing build (\d+)`)
					matches := buildRegex.FindSubmatch(buildSession.Out.Contents())
					buildID = string(matches[1])

					Eventually(buildSession).Should(gbytes.Say("waiting for /tmp/stop-waiting"))
				})

				AfterEach(func() {
					buildSession.Signal(os.Interrupt)
					<-buildSession.Exited
				})

				It("waits for the build", func() {
					Eventually(restartSession).Should(gbytes.Say(`Updating (instance|job)`))
					Consistently(restartSession, 5*time.Minute).ShouldNot(gexec.Exit())
				})

				It("finishes restarting once the build is done", func() {
					By("hijacking the build to tell it to finish")
					<-flyHijackTask(
						"-b", buildID,
						"-s", "one-off",
						"touch", "/tmp/stop-waiting",
					).Exited

					By("waiting for the build to exit")
					Eventually(buildSession).Should(gbytes.Say("done"))
					<-buildSession.Exited
					Expect(buildSession.ExitCode()).To(Equal(0))

					By("successfully restarting")
					<-restartSession.Exited
					Expect(restartSession.ExitCode()).To(Equal(0))
				})
			})
		})
	}

	Context("with one worker", func() {
		BeforeEach(func() {
			Deploy("deployments/one-forwarded-worker.yml")
			waitForRunningWorker()
		})

		describeRestartingTheWorker()
	})

	Context("with a single team worker", func() {
		BeforeEach(func() {
			Deploy("deployments/team-worker.yml")

			setTeam := spawnFlyInteractive(bytes.NewBufferString("y\n"), "set-team", "-n", "team-a", "--no-really-i-dont-want-any-auth")
			<-setTeam.Exited
			Expect(setTeam.ExitCode()).To(Equal(0))

			fly("login", "-c", atcExternalURL, "-n", "team-a")

			// wait for the team's worker to arrive now that team exists
			waitForRunningWorker()
		})

		describeRestartingTheWorker()
	})
})
<|endoftext|>


<commit_before>package physics

import (
	"math"
	"testing"

	gomock "gomock.googlecode.com/hg/gomock"

	. "chunkymonkey/types"
)

func almostEqual(v1, v2 float64) bool {
	return math.Fabs(v1-v2) < 1e-10
}

func Test_calcNextBlockDt(t *testing.T) {
	type Test struct {
		p        AbsCoord
		v        AbsVelocityCoord
		expected TickTime
	}

	tests := []Test{
		// Degenerate case of zero velocity
		Test{0.0, 0.0, math.MaxFloat64},
		Test{0.5, 0.0, math.MaxFloat64},
		Test{-0.5, 0.0, math.MaxFloat64},

		// Not sure if the apparent disparity between these two cases matters,
		// given the starting position is on a block boundary.
		Test{0.0, 0.1, 10.0},
		Test{0.0, -0.1, 0.0},

		// +ve pos, +ve vel
		Test{1.0, 0.1, 10.0},
		Test{1.0, 0.5, 2.0},
		Test{1.0, 1.0, 1.0},
		Test{20.0, 0.1, 10.0},
		Test{20.0, 0.5, 2.0},
		Test{20.0, 1.0, 1.0},

		// +ve pos, -ve vel
		Test{0.9, -0.1, 9.0},
		Test{0.9, -0.5, 1.8},
		Test{0.9, -1.0, 0.9},
		Test{19.9, -0.1, 9.0},
		Test{19.9, -0.5, 1.8},
		Test{19.9, -1.0, 0.9},

		// -ve pos, -ve vel
		Test{-1.0, -0.1, 10.0},
		Test{-1.0, -0.5, 2.0},
		Test{-1.0, -1.0, 1.0},
		Test{-20.0, -0.1, 10.0},
		Test{-20.0, -0.5, 2.0},
		Test{-20.0, -1.0, 1.0},

		// -ve pos, +ve vel
		Test{-0.9, 0.1, 9.0},
		Test{-0.9, 0.5, 1.8},
		Test{-0.9, 1.0, 0.9},
		Test{-19.9, 0.1, 9.0},
		Test{-19.9, 0.5, 1.8},
		Test{-19.9, 1.0, 0.9},

		// Crossing p=0
		Test{-0.5, 1.0, 0.5},
		Test{0.5, -1.0, 0.5},
	}

	for _, r := range tests {
		result := calcNextBlockDt(r.p, r.v)
		if !almostEqual(float64(r.expected), float64(result)) {
			t.Errorf("calcNextBlockDt(%g, %g) expected %g got %g",
				r.p, r.v, r.expected, result)
		}
	}
}

func Test_getBlockAxisMove(t *testing.T) {
	type Test struct {
		xDt, yDt, zDt    TickTime
		expBlockAxisMove blockAxisMove
		expDt            TickTime
	}

	tests := []Test{
		// In these first tests it doesn't really matter which axis is the
		// answer, but the code has a bias for the X axis so we roll with that.
		Test{0.0, 0.0, 0.0, blockAxisMoveX, 0.0},
		Test{0.5, 0.5, 0.5, blockAxisMoveX, 0.5},
		Test{1.0, 1.0, 1.0, blockAxisMoveX, 1.0},
		Test{1.5, 1.5, 1.5, blockAxisMoveX, 1.5},

		// Hit block boundary on X axis first
		Test{0.5, 0.9, 0.9, blockAxisMoveX, 0.5},
		Test{0.5, 10.0, 20.0, blockAxisMoveX, 0.5},
		Test{0.1, 20.0, 10.0, blockAxisMoveX, 0.1},

		// Hit block boundary on Y axis first
		Test{0.9, 0.5, 0.9, blockAxisMoveY, 0.5},
		Test{10.0, 0.5, 20.0, blockAxisMoveY, 0.5},
		Test{20.0, 0.1, 10.0, blockAxisMoveY, 0.1},

		// Hit block boundary on Z axis first
		Test{0.9, 0.9, 0.5, blockAxisMoveZ, 0.5},
		Test{10.0, 20.0, 0.5, blockAxisMoveZ, 0.5},
		Test{20.0, 10.0, 0.1, blockAxisMoveZ, 0.1},
	}

	for _, r := range tests {
		resultBlockAxisMove, resultDt := getBlockAxisMove(r.xDt, r.yDt, r.zDt)
		if r.expBlockAxisMove != resultBlockAxisMove || !almostEqual(float64(r.expDt), float64(resultDt)) {
			t.Errorf(
				"getBlockAxisMove(%g, %g, %g) expected %d,%g got %d,%g",
				r.xDt, r.yDt, r.zDt,
				r.expBlockAxisMove, r.expDt,
				resultBlockAxisMove, resultDt)
		}
	}
}

func Test_VelocityFromLook(t *testing.T) {
	type Test struct {
		look     LookDegrees
		momentum float64
		want     AbsVelocity
	}
	tests := []Test{}

	// TODO: Add some unit tests that accurately reflect the desired
	// behaviour.
	for _, test := range tests {
		v := VelocityFromLook(test.look, test.momentum)
		if v.X != test.want.X || v.Y != test.want.Y || v.Z != test.want.Z {
			t.Errorf("VelocityFromLook, wanted %+v, got %+v", test.want, v)
		}
	}
}

func tickFixtures(t *testing.T) (mockCtrl *gomock.Controller, mockBlockQuerier *MockIBlockQuerier, pointObj *PointObject) {
	mockCtrl = gomock.NewController(t)
	mockBlockQuerier = NewMockIBlockQuerier(mockCtrl)
	pointObj = new(PointObject)
	return
}

func Test_PointObject_Tick_FallsToImmediateSurface(t *testing.T) {
	mockCtrl, mockBlockQuerier, pointObj := tickFixtures(t)
	defer mockCtrl.Finish()

	pointObj.Init(
		&AbsXyz{0.5, 100.1, 0.5},
		&AbsVelocity{},
	)

	mockBlockQuerier.EXPECT().BlockQuery(BlockXyz{0, 99, 0}).Return(true, true)
	pointObj.Tick(mockBlockQuerier)

	expectedBlockPos := BlockXyz{0, 100, 0}
	if !pointObj.position.ToBlockXyz().Equals(expectedBlockPos) {
		t.Errorf("Expected object to end at %#v but was at %#v", expectedBlockPos, pointObj.position)
	}
}
<commit_msg>Added additional physics unit tests.<commit_after>package physics

import (
	"math"
	"testing"

	gomock "gomock.googlecode.com/hg/gomock"

	. "chunkymonkey/types"
)

func almostEqual(v1, v2 float64) bool {
	return math.Fabs(v1-v2) < 1e-10
}

func Test_calcNextBlockDt(t *testing.T) {
	type Test struct {
		p        AbsCoord
		v        AbsVelocityCoord
		expected TickTime
	}

	tests := []Test{
		// Degenerate case of zero velocity
		Test{0.0, 0.0, math.MaxFloat64},
		Test{0.5, 0.0, math.MaxFloat64},
		Test{-0.5, 0.0, math.MaxFloat64},

		// Not sure if the apparent disparity between these two cases matters,
		// given the starting position is on a block boundary.
		Test{0.0, 0.1, 10.0},
		Test{0.0, -0.1, 0.0},

		// +ve pos, +ve vel
		Test{1.0, 0.1, 10.0},
		Test{1.0, 0.5, 2.0},
		Test{1.0, 1.0, 1.0},
		Test{20.0, 0.1, 10.0},
		Test{20.0, 0.5, 2.0},
		Test{20.0, 1.0, 1.0},

		// +ve pos, -ve vel
		Test{0.9, -0.1, 9.0},
		Test{0.9, -0.5, 1.8},
		Test{0.9, -1.0, 0.9},
		Test{19.9, -0.1, 9.0},
		Test{19.9, -0.5, 1.8},
		Test{19.9, -1.0, 0.9},

		// -ve pos, -ve vel
		Test{-1.0, -0.1, 10.0},
		Test{-1.0, -0.5, 2.0},
		Test{-1.0, -1.0, 1.0},
		Test{-20.0, -0.1, 10.0},
		Test{-20.0, -0.5, 2.0},
		Test{-20.0, -1.0, 1.0},

		// -ve pos, +ve vel
		Test{-0.9, 0.1, 9.0},
		Test{-0.9, 0.5, 1.8},
		Test{-0.9, 1.0, 0.9},
		Test{-19.9, 0.1, 9.0},
		Test{-19.9, 0.5, 1.8},
		Test{-19.9, 1.0, 0.9},

		// Crossing p=0
		Test{-0.5, 1.0, 0.5},
		Test{0.5, -1.0, 0.5},
	}

	for _, r := range tests {
		result := calcNextBlockDt(r.p, r.v)
		if !almostEqual(float64(r.expected), float64(result)) {
			t.Errorf("calcNextBlockDt(%g, %g) expected %g got %g",
				r.p, r.v, r.expected, result)
		}
	}
}

func Test_getBlockAxisMove(t *testing.T) {
	type Test struct {
		xDt, yDt, zDt    TickTime
		expBlockAxisMove blockAxisMove
		expDt            TickTime
	}

	tests := []Test{
		// In these first tests it doesn't really matter which axis is the
		// answer, but the code has a bias for the X axis so we roll with that.
		Test{0.0, 0.0, 0.0, blockAxisMoveX, 0.0},
		Test{0.5, 0.5, 0.5, blockAxisMoveX, 0.5},
		Test{1.0, 1.0, 1.0, blockAxisMoveX, 1.0},
		Test{1.5, 1.5, 1.5, blockAxisMoveX, 1.5},

		// Hit block boundary on X axis first
		Test{0.5, 0.9, 0.9, blockAxisMoveX, 0.5},
		Test{0.5, 10.0, 20.0, blockAxisMoveX, 0.5},
		Test{0.1, 20.0, 10.0, blockAxisMoveX, 0.1},

		// Hit block boundary on Y axis first
		Test{0.9, 0.5, 0.9, blockAxisMoveY, 0.5},
		Test{10.0, 0.5, 20.0, blockAxisMoveY, 0.5},
		Test{20.0, 0.1, 10.0, blockAxisMoveY, 0.1},

		// Hit block boundary on Z axis first
		Test{0.9, 0.9, 0.5, blockAxisMoveZ, 0.5},
		Test{10.0, 20.0, 0.5, blockAxisMoveZ, 0.5},
		Test{20.0, 10.0, 0.1, blockAxisMoveZ, 0.1},
	}

	for _, r := range tests {
		resultBlockAxisMove, resultDt := getBlockAxisMove(r.xDt, r.yDt, r.zDt)
		if r.expBlockAxisMove != resultBlockAxisMove || !almostEqual(float64(r.expDt), float64(resultDt)) {
			t.Errorf(
				"getBlockAxisMove(%g, %g, %g) expected %d,%g got %d,%g",
				r.xDt, r.yDt, r.zDt,
				r.expBlockAxisMove, r.expDt,
				resultBlockAxisMove, resultDt)
		}
	}
}

func Test_VelocityFromLook(t *testing.T) {
	type Test struct {
		look     LookDegrees
		momentum float64
		want     AbsVelocity
	}
	tests := []Test{}

	// TODO: Add some unit tests that accurately reflect the desired
	// behaviour.
	for _, test := range tests {
		v := VelocityFromLook(test.look, test.momentum)
		if v.X != test.want.X || v.Y != test.want.Y || v.Z != test.want.Z {
			t.Errorf("VelocityFromLook, wanted %+v, got %+v", test.want, v)
		}
	}
}

func testTickFixtures(t *testing.T) (mockCtrl *gomock.Controller, mockBlockQuerier *MockIBlockQuerier, pointObj *PointObject) {
	mockCtrl = gomock.NewController(t)
	mockBlockQuerier = NewMockIBlockQuerier(mockCtrl)
	pointObj = new(PointObject)
	return
}

type test_PointObject_TickOnce struct {
	desc             string
	expectedEndBlock BlockXyz
	startPos         AbsXyz
	startVel         AbsVelocity
	queryBlock       BlockXyz
	isSolid          bool
}

func (test *test_PointObject_TickOnce) test(t *testing.T) {
	mockCtrl, mockBlockQuerier, pointObj := testTickFixtures(t)
	defer mockCtrl.Finish()

	t.Log(test.desc)

	pointObj.Init(&test.startPos, &test.startVel)

	mockBlockQuerier.EXPECT().BlockQuery(test.queryBlock).Return(test.isSolid, true)
	pointObj.Tick(mockBlockQuerier)

	if !pointObj.position.ToBlockXyz().Equals(test.expectedEndBlock) {
		t.Errorf("  Expected object to end at %#v but was at %#v", test.expectedEndBlock, pointObj.position)
	}
	t.Logf("  End position: %#v", pointObj.position)
	t.Logf("  End velocity: %#v", pointObj.velocity)
}

func Test_PointObject_TickOnce(t *testing.T) {
	tests := []test_PointObject_TickOnce{
		// Falling from rest:
		{
			desc:             "Fall to bottom of current block onto solid surface",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.1, 0.5},
			startVel:         AbsVelocity{},
			queryBlock:       BlockXyz{0, 99, 0},
			isSolid:          true,
		},
		// Moving with initial speed of 2 along an axis (from the middle of the
		// initial block this must hit or pass each side of the first block).
		{
			desc:             "Down to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{0, -2, 0},
			queryBlock:       BlockXyz{0, 99, 0},
			isSolid:          true,
		},
		{
			desc:             "Up to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{0, 2, 0},
			queryBlock:       BlockXyz{0, 101, 0},
			isSolid:          true,
		},
		{
			desc:             "North to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{-2, 0, 0},
			queryBlock:       BlockXyz{-1, 100, 0},
			isSolid:          true,
		},
		{
			desc:             "South to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{2, 0, 0},
			queryBlock:       BlockXyz{1, 100, 0},
			isSolid:          true,
		},
		{
			desc:             "East to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{0, 0, -2},
			queryBlock:       BlockXyz{0, 100, -1},
			isSolid:          true,
		},
		{
			desc:             "West to solid",
			expectedEndBlock: BlockXyz{0, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{0, 0, 2},
			queryBlock:       BlockXyz{0, 100, 1},
			isSolid:          true,
		},
		// Move with enough speed to move into another non-solid block but make no other transitions.
		{
			desc:             "Down to non-solid",
			expectedEndBlock: BlockXyz{0, 99, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{0, -1, 0},
			queryBlock:       BlockXyz{0, 99, 0},
			isSolid:          false,
		},
		{
			desc:             "North to non-solid",
			expectedEndBlock: BlockXyz{-1, 100, 0},
			startPos:         AbsXyz{0.5, 100.5, 0.5},
			startVel:         AbsVelocity{-1, 0, 0},
			queryBlock:       BlockXyz{-1, 100, 0},
			isSolid:          false,
		},
	}

	for i := range tests {
		test := &tests[i]
		test.test(t)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package analyze

import (
	"fmt"
	"strconv"

	"github.com/gyuho/dataframe"
)

// importBenchMetrics adds benchmark metrics from client-side
// and aggregates this to system metrics by unix timestamps.
func (data *analyzeData) importBenchMetrics(fpath string) (err error) {
	data.benchMetricsFilePath = fpath
	data.benchMetrics.frame, err = dataframe.NewFromCSV(nil, fpath)
	if err != nil {
		return
	}

	var unixTSColumn dataframe.Column
	unixTSColumn, err = data.benchMetrics.frame.Column("UNIX-TS")
	if err != nil {
		return err
	}

	// get first(minimum) unix second
	fv, ok := unixTSColumn.FrontNonNil()
	if !ok {
		return fmt.Errorf("FrontNonNil %s has empty Unix time %v", fpath, fv)
	}
	fs, ok := fv.String()
	if !ok {
		return fmt.Errorf("cannot String %v", fv)
	}
	data.benchMetrics.frontUnixTS, err = strconv.ParseInt(fs, 10, 64)
	if err != nil {
		return err
	}

	// get last(maximum) unix second
	bv, ok := unixTSColumn.BackNonNil()
	if !ok {
		return fmt.Errorf("BackNonNil %s has empty Unix time %v", fpath, fv)
	}
	bs, ok := bv.String()
	if !ok {
		return fmt.Errorf("cannot String %v", bv)
	}
	data.benchMetrics.lastUnixTS, err = strconv.ParseInt(bs, 10, 64)
	if err != nil {
		return err
	}

	return
}
<commit_msg>analyze: handle duplicate timestamps from benchmark<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package analyze

import (
	"fmt"
	"strconv"

	"github.com/gyuho/dataframe"
)

// importBenchMetrics adds benchmark metrics from client-side
// and aggregates this to system metrics by unix timestamps.
func (data *analyzeData) importBenchMetrics(fpath string) (err error) {
	data.benchMetricsFilePath = fpath
	data.benchMetrics.frame, err = dataframe.NewFromCSV(nil, fpath)
	if err != nil {
		return
	}

	var unixTSColumn dataframe.Column
	unixTSColumn, err = data.benchMetrics.frame.Column("UNIX-TS")
	if err != nil {
		return err
	}

	// get first(minimum) unix second
	fv, ok := unixTSColumn.FrontNonNil()
	if !ok {
		return fmt.Errorf("FrontNonNil %s has empty Unix time %v", fpath, fv)
	}
	fs, ok := fv.String()
	if !ok {
		return fmt.Errorf("cannot String %v", fv)
	}
	data.benchMetrics.frontUnixTS, err = strconv.ParseInt(fs, 10, 64)
	if err != nil {
		return err
	}

	nc := dataframe.NewColumn("UNIX-TS")
	for i := 0; i < unixTSColumn.Count(); i++ {
		nc.PushBack(dataframe.NewStringValue(fmt.Sprintf("%d", data.benchMetrics.frontUnixTS+int64(i))))
	}

	// get last(maximum) unix second
	bv, ok := nc.BackNonNil()
	if !ok {
		return fmt.Errorf("BackNonNil %s has empty Unix time %v", fpath, fv)
	}
	bs, ok := bv.String()
	if !ok {
		return fmt.Errorf("cannot String %v", bv)
	}
	data.benchMetrics.lastUnixTS, err = strconv.ParseInt(bs, 10, 64)
	if err != nil {
		return err
	}

	if ok = data.benchMetrics.frame.DeleteColumn("UNIX-TS"); !ok {
		return fmt.Errorf("UNIX-TS column is not deleted %v", data.benchMetrics.frame.Headers())
	}

	// overwrite duplicate timestamps
	if err = data.benchMetrics.frame.AddColumn(nc); err != nil {
		return err
	}
	if err = data.benchMetrics.frame.MoveColumn("UNIX-TS", 0); err != nil {
		return err
	}
	return
}
<|endoftext|>


<commit_before>package functions

import (
	"fmt"

	"github.com/influxdata/platform"
	"github.com/influxdata/platform/query"
	"github.com/influxdata/platform/query/execute"
	"github.com/influxdata/platform/query/functions/storage"
	"github.com/influxdata/platform/query/interpreter"
	"github.com/influxdata/platform/query/plan"
	"github.com/influxdata/platform/query/semantic"
	"github.com/pkg/errors"
)

const FromKind = "from"

type FromOpSpec struct {
	Database string      `json:"db,omitempty"`
	Bucket   string      `json:"bucket,omitempty"`
	BucketID platform.ID `json:"bucket_id,omitempty"`
	Hosts    []string    `json:"hosts"`
}

var fromSignature = semantic.FunctionSignature{
	Params: map[string]semantic.Type{
		"db": semantic.String,
	},
	ReturnType: query.TableObjectType,
}

func init() {
	query.RegisterFunction(FromKind, createFromOpSpec, fromSignature)
	query.RegisterOpSpec(FromKind, newFromOp)
	plan.RegisterProcedureSpec(FromKind, newFromProcedure, FromKind)
	execute.RegisterSource(FromKind, createFromSource)
}

func createFromOpSpec(args query.Arguments, a *query.Administration) (query.OperationSpec, error) {
	spec := new(FromOpSpec)

	if db, ok, err := args.GetString("db"); err != nil {
		return nil, err
	} else if ok {
		spec.Database = db
	}

	if bucket, ok, err := args.GetString("bucket"); err != nil {
		return nil, err
	} else if ok {
		spec.Bucket = bucket
	}

	if bucketID, ok, err := args.GetString("bucketID"); err != nil {
		return nil, err
	} else if ok {
		err := spec.BucketID.DecodeFromString(bucketID)
		if err != nil {
			return nil, errors.Wrap(err, "invalid bucket ID")
		}
	}

	if spec.Database == "" && spec.Bucket == "" && len(spec.BucketID) == 0 {
		return nil, errors.New("must specify one of db or bucket")
	}
	if spec.Database != "" && spec.Bucket != "" && len(spec.BucketID) == 0 {
		return nil, errors.New("must specify only one of db or bucket")
	}

	if array, ok, err := args.GetArray("hosts", semantic.String); err != nil {
		return nil, err
	} else if ok {
		spec.Hosts, err = interpreter.ToStringArray(array)
		if err != nil {
			return nil, err
		}
	}

	return spec, nil
}

func newFromOp() query.OperationSpec {
	return new(FromOpSpec)
}

func (s *FromOpSpec) Kind() query.OperationKind {
	return FromKind
}

type FromProcedureSpec struct {
	Database string
	Bucket   string
	BucketID platform.ID
	Hosts    []string

	BoundsSet bool
	Bounds    plan.BoundsSpec

	FilterSet bool
	Filter    *semantic.FunctionExpression

	DescendingSet bool
	Descending    bool

	LimitSet     bool
	PointsLimit  int64
	SeriesLimit  int64
	SeriesOffset int64

	WindowSet bool
	Window    plan.WindowSpec

	GroupingSet bool
	OrderByTime bool
	GroupMode   GroupMode
	GroupKeys   []string

	AggregateSet    bool
	AggregateMethod string
}

func newFromProcedure(qs query.OperationSpec, pa plan.Administration) (plan.ProcedureSpec, error) {
	spec, ok := qs.(*FromOpSpec)
	if !ok {
		return nil, fmt.Errorf("invalid spec type %T", qs)
	}

	return &FromProcedureSpec{
		Database: spec.Database,
		Bucket:   spec.Bucket,
		Hosts:    spec.Hosts,
	}, nil
}

func (s *FromProcedureSpec) Kind() plan.ProcedureKind {
	return FromKind
}
func (s *FromProcedureSpec) TimeBounds() plan.BoundsSpec {
	return s.Bounds
}
func (s *FromProcedureSpec) Copy() plan.ProcedureSpec {
	ns := new(FromProcedureSpec)

	ns.Database = s.Database
	ns.Bucket = s.Bucket
	if len(s.BucketID) > 0 {
		ns.BucketID = make(platform.ID, len(s.BucketID))
		copy(ns.BucketID, s.BucketID)
	}

	if len(s.Hosts) > 0 {
		ns.Hosts = make([]string, len(s.Hosts))
		copy(ns.Hosts, s.Hosts)
	}

	ns.BoundsSet = s.BoundsSet
	ns.Bounds = s.Bounds

	ns.FilterSet = s.FilterSet
	// TODO copy predicate
	ns.Filter = s.Filter

	ns.DescendingSet = s.DescendingSet
	ns.Descending = s.Descending

	ns.LimitSet = s.LimitSet
	ns.PointsLimit = s.PointsLimit
	ns.SeriesLimit = s.SeriesLimit
	ns.SeriesOffset = s.SeriesOffset

	ns.WindowSet = s.WindowSet
	ns.Window = s.Window

	ns.AggregateSet = s.AggregateSet
	ns.AggregateMethod = s.AggregateMethod

	return ns
}

func createFromSource(prSpec plan.ProcedureSpec, dsid execute.DatasetID, a execute.Administration) (execute.Source, error) {
	spec := prSpec.(*FromProcedureSpec)
	var w execute.Window
	if spec.WindowSet {
		w = execute.Window{
			Every:  execute.Duration(spec.Window.Every),
			Period: execute.Duration(spec.Window.Period),
			Round:  execute.Duration(spec.Window.Round),
			Start:  a.ResolveTime(spec.Window.Start),
		}
	} else {
		duration := execute.Duration(a.ResolveTime(spec.Bounds.Stop)) - execute.Duration(a.ResolveTime(spec.Bounds.Start))
		w = execute.Window{
			Every:  duration,
			Period: duration,
			Start:  a.ResolveTime(spec.Bounds.Start),
		}
	}
	currentTime := w.Start + execute.Time(w.Period)
	bounds := execute.Bounds{
		Start: a.ResolveTime(spec.Bounds.Start),
		Stop:  a.ResolveTime(spec.Bounds.Stop),
	}

	deps := a.Dependencies()[FromKind].(storage.Dependencies)
	orgID := a.OrganizationID()

	var bucketID platform.ID
	// Determine bucketID
	switch {
	case spec.Database != "":
		// The bucket ID will be treated as the database name
		bucketID = platform.ID(spec.Database)
	case len(spec.BucketID) != 0:
		bucketID = spec.BucketID
	case spec.Bucket != "":
		b, ok := deps.BucketLookup.Lookup(orgID, spec.Bucket)
		if !ok {
			return nil, fmt.Errorf("could not find bucket %q", spec.Bucket)
		}
		bucketID = b
	}

	return storage.NewSource(
		dsid,
		deps.Reader,
		storage.ReadSpec{
			OrganizationID:  orgID,
			BucketID:        bucketID,
			Hosts:           spec.Hosts,
			Predicate:       spec.Filter,
			PointsLimit:     spec.PointsLimit,
			SeriesLimit:     spec.SeriesLimit,
			SeriesOffset:    spec.SeriesOffset,
			Descending:      spec.Descending,
			OrderByTime:     spec.OrderByTime,
			GroupMode:       storage.GroupMode(spec.GroupMode),
			GroupKeys:       spec.GroupKeys,
			AggregateMethod: spec.AggregateMethod,
		},
		bounds,
		w,
		currentTime,
	), nil
}

func InjectFromDependencies(depsMap execute.Dependencies, deps storage.Dependencies) error {
	if err := deps.Validate(); err != nil {
		return err
	}
	depsMap[FromKind] = deps
	return nil
}
<commit_msg>fix(query): fix passing down the bucket id in the from call<commit_after>package functions

import (
	"fmt"

	"github.com/influxdata/platform"
	"github.com/influxdata/platform/query"
	"github.com/influxdata/platform/query/execute"
	"github.com/influxdata/platform/query/functions/storage"
	"github.com/influxdata/platform/query/interpreter"
	"github.com/influxdata/platform/query/plan"
	"github.com/influxdata/platform/query/semantic"
	"github.com/pkg/errors"
)

const FromKind = "from"

type FromOpSpec struct {
	Database string      `json:"db,omitempty"`
	Bucket   string      `json:"bucket,omitempty"`
	BucketID platform.ID `json:"bucket_id,omitempty"`
	Hosts    []string    `json:"hosts"`
}

var fromSignature = semantic.FunctionSignature{
	Params: map[string]semantic.Type{
		"db": semantic.String,
	},
	ReturnType: query.TableObjectType,
}

func init() {
	query.RegisterFunction(FromKind, createFromOpSpec, fromSignature)
	query.RegisterOpSpec(FromKind, newFromOp)
	plan.RegisterProcedureSpec(FromKind, newFromProcedure, FromKind)
	execute.RegisterSource(FromKind, createFromSource)
}

func createFromOpSpec(args query.Arguments, a *query.Administration) (query.OperationSpec, error) {
	spec := new(FromOpSpec)

	if db, ok, err := args.GetString("db"); err != nil {
		return nil, err
	} else if ok {
		spec.Database = db
	}

	if bucket, ok, err := args.GetString("bucket"); err != nil {
		return nil, err
	} else if ok {
		spec.Bucket = bucket
	}

	if bucketID, ok, err := args.GetString("bucketID"); err != nil {
		return nil, err
	} else if ok {
		err := spec.BucketID.DecodeFromString(bucketID)
		if err != nil {
			return nil, errors.Wrap(err, "invalid bucket ID")
		}
	}

	if spec.Database == "" && spec.Bucket == "" && len(spec.BucketID) == 0 {
		return nil, errors.New("must specify one of db or bucket")
	}
	if spec.Database != "" && spec.Bucket != "" && len(spec.BucketID) == 0 {
		return nil, errors.New("must specify only one of db or bucket")
	}

	if array, ok, err := args.GetArray("hosts", semantic.String); err != nil {
		return nil, err
	} else if ok {
		spec.Hosts, err = interpreter.ToStringArray(array)
		if err != nil {
			return nil, err
		}
	}

	return spec, nil
}

func newFromOp() query.OperationSpec {
	return new(FromOpSpec)
}

func (s *FromOpSpec) Kind() query.OperationKind {
	return FromKind
}

type FromProcedureSpec struct {
	Database string
	Bucket   string
	BucketID platform.ID
	Hosts    []string

	BoundsSet bool
	Bounds    plan.BoundsSpec

	FilterSet bool
	Filter    *semantic.FunctionExpression

	DescendingSet bool
	Descending    bool

	LimitSet     bool
	PointsLimit  int64
	SeriesLimit  int64
	SeriesOffset int64

	WindowSet bool
	Window    plan.WindowSpec

	GroupingSet bool
	OrderByTime bool
	GroupMode   GroupMode
	GroupKeys   []string

	AggregateSet    bool
	AggregateMethod string
}

func newFromProcedure(qs query.OperationSpec, pa plan.Administration) (plan.ProcedureSpec, error) {
	spec, ok := qs.(*FromOpSpec)
	if !ok {
		return nil, fmt.Errorf("invalid spec type %T", qs)
	}

	return &FromProcedureSpec{
		Database: spec.Database,
		Bucket:   spec.Bucket,
		BucketID: spec.BucketID,
		Hosts:    spec.Hosts,
	}, nil
}

func (s *FromProcedureSpec) Kind() plan.ProcedureKind {
	return FromKind
}
func (s *FromProcedureSpec) TimeBounds() plan.BoundsSpec {
	return s.Bounds
}
func (s *FromProcedureSpec) Copy() plan.ProcedureSpec {
	ns := new(FromProcedureSpec)

	ns.Database = s.Database
	ns.Bucket = s.Bucket
	if len(s.BucketID) > 0 {
		ns.BucketID = make(platform.ID, len(s.BucketID))
		copy(ns.BucketID, s.BucketID)
	}

	if len(s.Hosts) > 0 {
		ns.Hosts = make([]string, len(s.Hosts))
		copy(ns.Hosts, s.Hosts)
	}

	ns.BoundsSet = s.BoundsSet
	ns.Bounds = s.Bounds

	ns.FilterSet = s.FilterSet
	// TODO copy predicate
	ns.Filter = s.Filter

	ns.DescendingSet = s.DescendingSet
	ns.Descending = s.Descending

	ns.LimitSet = s.LimitSet
	ns.PointsLimit = s.PointsLimit
	ns.SeriesLimit = s.SeriesLimit
	ns.SeriesOffset = s.SeriesOffset

	ns.WindowSet = s.WindowSet
	ns.Window = s.Window

	ns.AggregateSet = s.AggregateSet
	ns.AggregateMethod = s.AggregateMethod

	return ns
}

func createFromSource(prSpec plan.ProcedureSpec, dsid execute.DatasetID, a execute.Administration) (execute.Source, error) {
	spec := prSpec.(*FromProcedureSpec)
	var w execute.Window
	if spec.WindowSet {
		w = execute.Window{
			Every:  execute.Duration(spec.Window.Every),
			Period: execute.Duration(spec.Window.Period),
			Round:  execute.Duration(spec.Window.Round),
			Start:  a.ResolveTime(spec.Window.Start),
		}
	} else {
		duration := execute.Duration(a.ResolveTime(spec.Bounds.Stop)) - execute.Duration(a.ResolveTime(spec.Bounds.Start))
		w = execute.Window{
			Every:  duration,
			Period: duration,
			Start:  a.ResolveTime(spec.Bounds.Start),
		}
	}
	currentTime := w.Start + execute.Time(w.Period)
	bounds := execute.Bounds{
		Start: a.ResolveTime(spec.Bounds.Start),
		Stop:  a.ResolveTime(spec.Bounds.Stop),
	}

	deps := a.Dependencies()[FromKind].(storage.Dependencies)
	orgID := a.OrganizationID()

	var bucketID platform.ID
	// Determine bucketID
	switch {
	case spec.Database != "":
		// The bucket ID will be treated as the database name
		bucketID = platform.ID(spec.Database)
	case len(spec.BucketID) != 0:
		bucketID = spec.BucketID
	case spec.Bucket != "":
		b, ok := deps.BucketLookup.Lookup(orgID, spec.Bucket)
		if !ok {
			return nil, fmt.Errorf("could not find bucket %q", spec.Bucket)
		}
		bucketID = b
	}

	return storage.NewSource(
		dsid,
		deps.Reader,
		storage.ReadSpec{
			OrganizationID:  orgID,
			BucketID:        bucketID,
			Hosts:           spec.Hosts,
			Predicate:       spec.Filter,
			PointsLimit:     spec.PointsLimit,
			SeriesLimit:     spec.SeriesLimit,
			SeriesOffset:    spec.SeriesOffset,
			Descending:      spec.Descending,
			OrderByTime:     spec.OrderByTime,
			GroupMode:       storage.GroupMode(spec.GroupMode),
			GroupKeys:       spec.GroupKeys,
			AggregateMethod: spec.AggregateMethod,
		},
		bounds,
		w,
		currentTime,
	), nil
}

func InjectFromDependencies(depsMap execute.Dependencies, deps storage.Dependencies) error {
	if err := deps.Validate(); err != nil {
		return err
	}
	depsMap[FromKind] = deps
	return nil
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/service/xray"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSXraySamplingRule_basic(t *testing.T) {
	var samplingRule xray.SamplingRule
	resourceName := "aws_xray_sampling_rule.test"
	rString := acctest.RandString(8)
	ruleName := fmt.Sprintf("tf_acc_sampling_rule_%s", rString)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSXraySamplingRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSXraySamplingRuleConfig_basic(ruleName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttr(resourceName, "priority", "5"),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "reservoir_size", "10"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "1"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccAWSXraySamplingRule_update(t *testing.T) {
	var samplingRule xray.SamplingRule
	resourceName := "aws_xray_sampling_rule.test"
	rString := acctest.RandString(8)
	ruleName := fmt.Sprintf("tf_acc_sampling_rule_%s", rString)
	updatedPriority := acctest.RandIntRange(0, 9999)
	updatedReservoirSize := acctest.RandIntRange(0, 2147483647)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSXraySamplingRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSXraySamplingRuleConfig_update(ruleName, acctest.RandIntRange(0, 9999), acctest.RandIntRange(0, 2147483647)),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttrSet(resourceName, "priority"),
					resource.TestCheckResourceAttrSet(resourceName, "reservoir_size"),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "0"),
				),
			},
			{ // Update attributes
				Config: testAccAWSXraySamplingRuleConfig_update(ruleName, updatedPriority, updatedReservoirSize),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttr(resourceName, "priority", fmt.Sprintf("%d", updatedPriority)),
					resource.TestCheckResourceAttr(resourceName, "reservoir_size", fmt.Sprintf("%d", updatedReservoirSize)),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "0"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccCheckXraySamplingRuleExists(n string, samplingRule *xray.SamplingRule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No XRay Sampling Rule ID is set")
		}
		conn := testAccProvider.Meta().(*AWSClient).xrayconn

		rule, err := getXraySamplingRule(conn, rs.Primary.ID)

		if err != nil {
			return err
		}

		*samplingRule = *rule

		return nil
	}
}

func testAccCheckAWSXraySamplingRuleDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_xray_sampling_rule" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).xrayconn

		rule, err := getXraySamplingRule(conn, rs.Primary.ID)

		if err != nil {
			return err
		}

		if rule != nil {
			return fmt.Errorf("Expected XRay Sampling Rule to be destroyed, %s found", rs.Primary.ID)
		}
	}

	return nil
}

func testAccAWSXraySamplingRuleConfig_basic(ruleName string) string {
	return fmt.Sprintf(`
resource "aws_xray_sampling_rule" "test" {
	rule_name = "%s"
	priority = 5
	reservoir_size = 10
	url_path = "*"
	host = "*"
	http_method = "GET"
	service_type = "*"
	service_name = "*"
	fixed_rate = 0.3
	resource_arn = "*"
	version = 1
	attributes = {
		Hello = "World"
	}
}
`, ruleName)
}

func testAccAWSXraySamplingRuleConfig_update(ruleName string, priority int, reservoirSize int) string {
	return fmt.Sprintf(`
resource "aws_xray_sampling_rule" "test" {
	rule_name = "%s"
	priority = %d
	reservoir_size = %d
	url_path = "*"
	host = "*"
	http_method = "GET"
	service_type = "*"
	service_name = "*"
	fixed_rate = 0.3
	resource_arn = "*"
	version = 1
}
`, ruleName, priority, reservoirSize)
}
<commit_msg>tests/service/xray: Add PreCheck for service availability<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/service/xray"
	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSXraySamplingRule_basic(t *testing.T) {
	var samplingRule xray.SamplingRule
	resourceName := "aws_xray_sampling_rule.test"
	rString := acctest.RandString(8)
	ruleName := fmt.Sprintf("tf_acc_sampling_rule_%s", rString)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t); testAccPreCheckAWSXray(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSXraySamplingRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSXraySamplingRuleConfig_basic(ruleName),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttr(resourceName, "priority", "5"),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "reservoir_size", "10"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "1"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func TestAccAWSXraySamplingRule_update(t *testing.T) {
	var samplingRule xray.SamplingRule
	resourceName := "aws_xray_sampling_rule.test"
	rString := acctest.RandString(8)
	ruleName := fmt.Sprintf("tf_acc_sampling_rule_%s", rString)
	updatedPriority := acctest.RandIntRange(0, 9999)
	updatedReservoirSize := acctest.RandIntRange(0, 2147483647)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t); testAccPreCheckAWSXray(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSXraySamplingRuleDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAWSXraySamplingRuleConfig_update(ruleName, acctest.RandIntRange(0, 9999), acctest.RandIntRange(0, 2147483647)),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttrSet(resourceName, "priority"),
					resource.TestCheckResourceAttrSet(resourceName, "reservoir_size"),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "0"),
				),
			},
			{ // Update attributes
				Config: testAccAWSXraySamplingRuleConfig_update(ruleName, updatedPriority, updatedReservoirSize),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckXraySamplingRuleExists(resourceName, &samplingRule),
					testAccCheckResourceAttrRegionalARN(resourceName, "arn", "xray", fmt.Sprintf("sampling-rule/%s", ruleName)),
					resource.TestCheckResourceAttr(resourceName, "priority", fmt.Sprintf("%d", updatedPriority)),
					resource.TestCheckResourceAttr(resourceName, "reservoir_size", fmt.Sprintf("%d", updatedReservoirSize)),
					resource.TestCheckResourceAttr(resourceName, "version", "1"),
					resource.TestCheckResourceAttr(resourceName, "url_path", "*"),
					resource.TestCheckResourceAttr(resourceName, "host", "*"),
					resource.TestCheckResourceAttr(resourceName, "http_method", "GET"),
					resource.TestCheckResourceAttr(resourceName, "fixed_rate", "0.3"),
					resource.TestCheckResourceAttr(resourceName, "resource_arn", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_name", "*"),
					resource.TestCheckResourceAttr(resourceName, "service_type", "*"),
					resource.TestCheckResourceAttr(resourceName, "attributes.%", "0"),
				),
			},
			{
				ResourceName:      resourceName,
				ImportState:       true,
				ImportStateVerify: true,
			},
		},
	})
}

func testAccCheckXraySamplingRuleExists(n string, samplingRule *xray.SamplingRule) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[n]
		if !ok {
			return fmt.Errorf("Not found: %s", n)
		}

		if rs.Primary.ID == "" {
			return fmt.Errorf("No XRay Sampling Rule ID is set")
		}
		conn := testAccProvider.Meta().(*AWSClient).xrayconn

		rule, err := getXraySamplingRule(conn, rs.Primary.ID)

		if err != nil {
			return err
		}

		*samplingRule = *rule

		return nil
	}
}

func testAccCheckAWSXraySamplingRuleDestroy(s *terraform.State) error {
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_xray_sampling_rule" {
			continue
		}

		conn := testAccProvider.Meta().(*AWSClient).xrayconn

		rule, err := getXraySamplingRule(conn, rs.Primary.ID)

		if err != nil {
			return err
		}

		if rule != nil {
			return fmt.Errorf("Expected XRay Sampling Rule to be destroyed, %s found", rs.Primary.ID)
		}
	}

	return nil
}

func testAccPreCheckAWSXray(t *testing.T) {
	conn := testAccProvider.Meta().(*AWSClient).xrayconn

	input := &xray.GetSamplingRulesInput{}

	_, err := conn.GetSamplingRules(input)

	if testAccPreCheckSkipError(err) {
		t.Skipf("skipping acceptance testing: %s", err)
	}

	if err != nil {
		t.Fatalf("unexpected PreCheck error: %s", err)
	}
}

func testAccAWSXraySamplingRuleConfig_basic(ruleName string) string {
	return fmt.Sprintf(`
resource "aws_xray_sampling_rule" "test" {
	rule_name = "%s"
	priority = 5
	reservoir_size = 10
	url_path = "*"
	host = "*"
	http_method = "GET"
	service_type = "*"
	service_name = "*"
	fixed_rate = 0.3
	resource_arn = "*"
	version = 1
	attributes = {
		Hello = "World"
	}
}
`, ruleName)
}

func testAccAWSXraySamplingRuleConfig_update(ruleName string, priority int, reservoirSize int) string {
	return fmt.Sprintf(`
resource "aws_xray_sampling_rule" "test" {
	rule_name = "%s"
	priority = %d
	reservoir_size = %d
	url_path = "*"
	host = "*"
	http_method = "GET"
	service_type = "*"
	service_name = "*"
	fixed_rate = 0.3
	resource_arn = "*"
	version = 1
}
`, ruleName, priority, reservoirSize)
}
<|endoftext|>