<commit_before>/*
Copyright © 2021 Google

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 cmd

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/hashicorp/go-retryablehttp"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"golang.org/x/oauth2/google"
	computealpha "google.golang.org/api/compute/v0.alpha"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/container/v1"
	"google.golang.org/api/option"
	"legacymigration/pkg"
	"legacymigration/pkg/clusters"
	"legacymigration/pkg/migrate"
	"legacymigration/pkg/networks"
	"legacymigration/pkg/operations"
)

const (
	// Flag name constants.
	projectFlag                    = "project"
	containerBasePathFlag          = "container-base-url"
	networkFlag                    = "network"
	concurrentClustersFlag         = "concurrent-clusters"
	desiredControlPlaneVersionFlag = "control-plane-version"
	desiredNodeVersionFlag         = "node-version"
	pollingIntervalFlag            = "polling-interval"
	pollingDeadlineFlag            = "polling-deadline"
	inPlaceControlPlaneUpgradeFlag = "in-place-control-plane"
	validateOnlyFlag               = "validate-only"

	ConcurrentNetworks  = 1
	ConcurrentNodePools = 1
)

type migrateOptions struct {
	// Options set by flags.
	projectID                  string
	containerBasePath          string
	selectedNetwork            string
	concurrentClusters         uint16
	desiredControlPlaneVersion string
	desiredNodeVersion         string
	inPlaceControlPlaneUpgrade bool
	validateOnly               bool
	pollingInterval            time.Duration
	pollingDeadline            time.Duration

	// Field used for faking clients during tests.
	fetchClientFunc func(ctx context.Context, basePath string, authedClient *http.Client) (*pkg.Clients, error)

	// Options set during Complete
	clients   *pkg.Clients
	migrators []migrate.Migrator
}

var (
	rootCmd = newRootCmd()
)

// rootCmd represents the base command when called without any subcommands
func newRootCmd() *cobra.Command {
	o := migrateOptions{
		fetchClientFunc: fetchClients,
	}
	ctx, cancel := context.WithCancel(context.Background())

	cmd := &cobra.Command{
		Use:   "gkeconvert",
		Short: "Convert a GCE legacy network to a VPC network and upgrade GKE clusters.",
		Long: `This script converts a GCE legacy network to a VPC network (by switching
the network to custom subnet mode). It then performs GKE cluster upgrades to ensure
the clusters are compatible with a VPC network.`,

		PreRun: func(cmd *cobra.Command, args []string) {
			cobra.CheckErr(o.ValidateFlags())
			setupCloseHandler(cancel)
		},
		Run: func(cmd *cobra.Command, args []string) {
			cobra.CheckErr(o.Complete(ctx))
			cobra.CheckErr(o.Run(ctx))
		},
	}

	flags := cmd.Flags()

	flags.StringVarP(&o.projectID, projectFlag, "p", o.projectID, "project ID")

	// Target network options.
	flags.StringVarP(&o.selectedNetwork, networkFlag, "n", o.selectedNetwork, "GCE network to process.")

	// Concurrency options.
	flags.Uint16VarP(&o.concurrentClusters, concurrentClustersFlag, "C", 1, "Number of clusters per network to upgrade concurrently.")

	// Polling options.
	flags.DurationVar(&o.pollingInterval, pollingIntervalFlag, 10*time.Second, "Period between polling attempts.")
	flags.DurationVar(&o.pollingDeadline, pollingDeadlineFlag, 1*time.Hour, "Deadline for a long running operation to complete (e.g. to upgrade the cluster control plane).")

	// Cluster upgrade options.
	flags.StringVar(&o.desiredControlPlaneVersion, desiredControlPlaneVersionFlag, o.desiredControlPlaneVersion,
		`Desired GKE version for all cluster control planes.
For more information, see https://cloud.google.com/kubernetes-engine/versioning#versioning_scheme
Note:
  This version must be equal to or greater than the lowest control plane version on the network.`)
	flags.StringVar(&o.desiredNodeVersion, desiredNodeVersionFlag, o.desiredNodeVersion,
		`Desired GKE version for all cluster nodes. For more information, see https://cloud.google.com/kubernetes-engine/versioning#versioning_scheme
Note:
  This version must be greater than the lowest node pool version on the network as node pools cannot be upgraded in-place.`)

	flags.BoolVar(&o.inPlaceControlPlaneUpgrade, inPlaceControlPlaneUpgradeFlag, false,
		`Perform in-place control plane upgrade for all clusters.`)

	flags.BoolVar(&o.validateOnly, validateOnlyFlag, true,
		`Only run validation on the network and cluster resources; do not perform conversion`)

	// Test options.
	flags.StringVar(&o.containerBasePath, containerBasePathFlag, o.containerBasePath, "Custom URL for the container API endpoint (for testing).")

	cmd.MarkFlagRequired(projectFlag)
	cmd.MarkFlagRequired(networkFlag)
	cmd.MarkFlagRequired(desiredNodeVersionFlag)
	flags.MarkHidden(containerBasePathFlag)

	return cmd
}

// Execute runs the root command.
func Execute() {
	cobra.CheckErr(rootCmd.Execute())
}

// ValidateFlags ensures flags values are valid for execution.
func (o *migrateOptions) ValidateFlags() error {
	if o.projectID == "" {
		return fmt.Errorf("--%s not provided or empty", projectFlag)
	}
	if o.selectedNetwork == "" {
		return fmt.Errorf("--%s not provided or empty", networkFlag)
	}

	// Concurrency validation.
	if o.concurrentClusters < 1 {
		return fmt.Errorf("--%s must be an integer greater than 0", concurrentClustersFlag)
	}

	// Polling validation.
	if o.pollingInterval < 10*time.Second {
		return fmt.Errorf("--%s must greater than or equal to 10 seconds. Note: Upgrade operations times are O(minutes)", pollingIntervalFlag)
	}
	if o.pollingDeadline < 5*time.Minute {
		return fmt.Errorf("--%s must greater than or equal to 5 minutes. Note: Upgrade operations times are O(minutes)", pollingDeadlineFlag)
	}
	if o.pollingInterval > o.pollingDeadline {
		return fmt.Errorf("--%s=%v must be greater than --%s=%v", pollingDeadlineFlag, o.pollingDeadline, pollingIntervalFlag, o.pollingInterval)
	}

	// Version validation.
	if (o.desiredControlPlaneVersion == "" && !o.inPlaceControlPlaneUpgrade) ||
		(o.desiredControlPlaneVersion != "" && o.inPlaceControlPlaneUpgrade) {
		return fmt.Errorf("specify --%s or provide a version for --%s, but not both", inPlaceControlPlaneUpgradeFlag, desiredControlPlaneVersionFlag)
	}
	if o.desiredControlPlaneVersion != "" {
		if err := clusters.IsFormatValid(o.desiredControlPlaneVersion); err != nil {
			return fmt.Errorf("--%s=%q is not valid: %v", desiredControlPlaneVersionFlag, o.desiredControlPlaneVersion, err)
		}
	}
	if err := clusters.IsFormatValid(o.desiredNodeVersion); err != nil {
		return fmt.Errorf("--%s=%q is not valid: %v", desiredNodeVersionFlag, o.desiredNodeVersion, err)
	}
	// Use of `-` or `latest` aliases are validated later at the control plane and node pool level.
	if !o.inPlaceControlPlaneUpgrade &&
		o.desiredControlPlaneVersion != clusters.DefaultVersion &&
		o.desiredControlPlaneVersion != clusters.LatestVersion &&
		o.desiredNodeVersion != clusters.DefaultVersion &&
		o.desiredNodeVersion != clusters.LatestVersion {
		if err := clusters.IsWithinVersionSkew(o.desiredNodeVersion, o.desiredControlPlaneVersion, clusters.MaxVersionSkew); err != nil {
			return err
		}
	}

	return nil
}

// Complete cascades down the resource hierarchy, ensuring that all descendent migrators are initialized.
func (o *migrateOptions) Complete(ctx context.Context) error {
	authedClient, err := google.DefaultClient(ctx, compute.CloudPlatformScope)
	if err != nil {
		return err
	}

	o.clients, err = o.fetchClientFunc(ctx, o.containerBasePath, authedClient)
	if err != nil {
		return err
	}

	handler := operations.NewHandler(o.pollingInterval, o.pollingDeadline)
	options := &clusters.Options{
		ConcurrentNodePools:        ConcurrentNodePools,
		DesiredControlPlaneVersion: o.desiredControlPlaneVersion,
		DesiredNodeVersion:         o.desiredNodeVersion,
		InPlaceControlPlaneUpgrade: o.inPlaceControlPlaneUpgrade,
	}

	factory := func(n *compute.Network) migrate.Migrator {
		return networks.New(o.projectID, n, handler, o.clients, o.concurrentClusters, options)
	}

	log.Infof("Fetching network %s for project %q", o.selectedNetwork, o.projectID)

	ns, err := o.clients.Compute.ListNetworks(ctx, o.projectID)
	if err != nil {
		return fmt.Errorf("error listing networks: %w", err)
	}

	o.migrators = make([]migrate.Migrator, 0)
	for _, n := range ns {
		if n.Name == o.selectedNetwork {
			o.migrators = append(o.migrators, factory(n))
		}
	}

	if len(o.migrators) == 0 {
		return fmt.Errorf("unable to find network %s", o.selectedNetwork)
	}

	return nil
}

// Run cascades down the resource hierarchy, initializing, validating, and converting all descendent migrators.
func (o *migrateOptions) Run(ctx context.Context) error {
	sem := make(chan struct{}, ConcurrentNetworks)

	log.Info("Initialize objects for conversion.")
	if err := migrate.Complete(ctx, sem, o.migrators...); err != nil {
		return err
	}

	log.Info("Validate resources for conversion.")
	if err := migrate.Validate(ctx, sem, o.migrators...); err != nil {
		return err
	}

	if o.validateOnly {
		log.Infof("--%s==true; skipping conversion.", validateOnlyFlag)
		return nil
	}

	log.Info("Initiate resource conversion.")
	return migrate.Migrate(ctx, sem, o.migrators...)
}

// setupCloseHandler cancels the shared context when the user hits ctrl-c.
func setupCloseHandler(cancel context.CancelFunc) {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		cancel()
	}()
}

func getRetryableClientOption(retry int, waitMin, waitMax time.Duration, authedClient *http.Client) option.ClientOption {
	retryClient := retryablehttp.NewClient()
	retryClient.RetryMax = retry
	retryClient.RetryWaitMin = waitMin
	retryClient.RetryWaitMax = waitMax
	retryClient.Logger = nil
	retryClient.CheckRetry = retryPolicy()

	c := retryClient.StandardClient()
	c.Transport.(*retryablehttp.RoundTripper).Client.HTTPClient = authedClient
	return option.WithHTTPClient(c)
}

func fetchClients(ctx context.Context, basePath string, authedClient *http.Client) (*pkg.Clients, error) {
	opt := getRetryableClientOption(3, 5*time.Second, 30*time.Second, authedClient)
	computeService, err := compute.NewService(ctx, opt)
	if err != nil {
		return nil, err
	}
	containerService, err := container.NewService(ctx, opt)
	if err != nil {
		return nil, err
	}

	// Retry for up-to 5 minutes for Compute Alpha API calls.
	alphaOpt := getRetryableClientOption(5, 5*time.Second, 160*time.Second, authedClient)
	computeServiceAlpha, err := computealpha.NewService(ctx, alphaOpt)
	if err != nil {
		return nil, err
	}

	if basePath != "" {
		containerService.BasePath = basePath
	}
	return &pkg.Clients{
		Compute: &pkg.Compute{
			V1:    computeService,
			Alpha: computeServiceAlpha,
		},
		Container: &pkg.Container{V1: containerService},
	}, nil
}

func retryPolicy() func(ctx context.Context, resp *http.Response, err error) (bool, error) {
	return func(ctx context.Context, resp *http.Response, err error) (bool, error) {
		shouldRetry, newErr := retryablehttp.DefaultRetryPolicy(ctx, resp, err)
		if newErr != nil || shouldRetry == true {
			return shouldRetry, newErr
		}

		// GCE returns 403 as a RateLimiting response code.
		if resp.StatusCode == 403 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return true, err
			}
			return true, errors.New(string(body))
		}
		return false, nil
	}
}
<commit_msg>Increase default values for polling options.<commit_after>/*
Copyright © 2021 Google

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 cmd

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"legacymigration/pkg"
	"legacymigration/pkg/clusters"
	"legacymigration/pkg/migrate"
	"legacymigration/pkg/networks"
	"legacymigration/pkg/operations"

	"github.com/hashicorp/go-retryablehttp"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"golang.org/x/oauth2/google"
	computealpha "google.golang.org/api/compute/v0.alpha"
	"google.golang.org/api/compute/v1"
	"google.golang.org/api/container/v1"
	"google.golang.org/api/option"
)

const (
	// Flag name constants.
	projectFlag                    = "project"
	containerBasePathFlag          = "container-base-url"
	networkFlag                    = "network"
	concurrentClustersFlag         = "concurrent-clusters"
	desiredControlPlaneVersionFlag = "control-plane-version"
	desiredNodeVersionFlag         = "node-version"
	pollingIntervalFlag            = "polling-interval"
	pollingDeadlineFlag            = "polling-deadline"
	inPlaceControlPlaneUpgradeFlag = "in-place-control-plane"
	validateOnlyFlag               = "validate-only"

	ConcurrentNetworks  = 1
	ConcurrentNodePools = 1
)

type migrateOptions struct {
	// Options set by flags.
	projectID                  string
	containerBasePath          string
	selectedNetwork            string
	concurrentClusters         uint16
	desiredControlPlaneVersion string
	desiredNodeVersion         string
	inPlaceControlPlaneUpgrade bool
	validateOnly               bool
	pollingInterval            time.Duration
	pollingDeadline            time.Duration

	// Field used for faking clients during tests.
	fetchClientFunc func(ctx context.Context, basePath string, authedClient *http.Client) (*pkg.Clients, error)

	// Options set during Complete
	clients   *pkg.Clients
	migrators []migrate.Migrator
}

var (
	rootCmd = newRootCmd()
)

// rootCmd represents the base command when called without any subcommands
func newRootCmd() *cobra.Command {
	o := migrateOptions{
		fetchClientFunc: fetchClients,
	}
	ctx, cancel := context.WithCancel(context.Background())

	cmd := &cobra.Command{
		Use:   "gkeconvert",
		Short: "Convert a GCE legacy network to a VPC network and upgrade GKE clusters.",
		Long: `This script converts a GCE legacy network to a VPC network (by switching
the network to custom subnet mode). It then performs GKE cluster upgrades to ensure
the clusters are compatible with a VPC network.`,

		PreRun: func(cmd *cobra.Command, args []string) {
			cobra.CheckErr(o.ValidateFlags())
			setupCloseHandler(cancel)
		},
		Run: func(cmd *cobra.Command, args []string) {
			cobra.CheckErr(o.Complete(ctx))
			cobra.CheckErr(o.Run(ctx))
		},
	}

	flags := cmd.Flags()

	flags.StringVarP(&o.projectID, projectFlag, "p", o.projectID, "project ID")

	// Target network options.
	flags.StringVarP(&o.selectedNetwork, networkFlag, "n", o.selectedNetwork, "GCE network to process.")

	// Concurrency options.
	flags.Uint16VarP(&o.concurrentClusters, concurrentClustersFlag, "C", 1, "Number of clusters per network to upgrade concurrently.")

	// Polling options.
	flags.DurationVar(&o.pollingInterval, pollingIntervalFlag, 15*time.Second, "Period between polling attempts.")
	flags.DurationVar(&o.pollingDeadline, pollingDeadlineFlag, 24*time.Hour, "Deadline for a long running operation to complete (e.g. to upgrade a cluster node pool).")

	// Cluster upgrade options.
	flags.StringVar(&o.desiredControlPlaneVersion, desiredControlPlaneVersionFlag, o.desiredControlPlaneVersion,
		`Desired GKE version for all cluster control planes.
For more information, see https://cloud.google.com/kubernetes-engine/versioning#versioning_scheme
Note:
  This version must be equal to or greater than the lowest control plane version on the network.`)
	flags.StringVar(&o.desiredNodeVersion, desiredNodeVersionFlag, o.desiredNodeVersion,
		`Desired GKE version for all cluster nodes. For more information, see https://cloud.google.com/kubernetes-engine/versioning#versioning_scheme
Note:
  This version must be greater than the lowest node pool version on the network as node pools cannot be upgraded in-place.`)

	flags.BoolVar(&o.inPlaceControlPlaneUpgrade, inPlaceControlPlaneUpgradeFlag, false,
		`Perform in-place control plane upgrade for all clusters.`)

	flags.BoolVar(&o.validateOnly, validateOnlyFlag, true,
		`Only run validation on the network and cluster resources; do not perform conversion`)

	// Test options.
	flags.StringVar(&o.containerBasePath, containerBasePathFlag, o.containerBasePath, "Custom URL for the container API endpoint (for testing).")

	cmd.MarkFlagRequired(projectFlag)
	cmd.MarkFlagRequired(networkFlag)
	cmd.MarkFlagRequired(desiredNodeVersionFlag)
	flags.MarkHidden(containerBasePathFlag)

	return cmd
}

// Execute runs the root command.
func Execute() {
	cobra.CheckErr(rootCmd.Execute())
}

// ValidateFlags ensures flags values are valid for execution.
func (o *migrateOptions) ValidateFlags() error {
	if o.projectID == "" {
		return fmt.Errorf("--%s not provided or empty", projectFlag)
	}
	if o.selectedNetwork == "" {
		return fmt.Errorf("--%s not provided or empty", networkFlag)
	}

	// Concurrency validation.
	if o.concurrentClusters < 1 {
		return fmt.Errorf("--%s must be an integer greater than 0", concurrentClustersFlag)
	}

	// Polling validation.
	if o.pollingInterval < 10*time.Second {
		return fmt.Errorf("--%s must greater than or equal to 10 seconds. Note: Upgrade operations times are O(minutes)", pollingIntervalFlag)
	}
	if o.pollingDeadline < 5*time.Minute {
		return fmt.Errorf("--%s must greater than or equal to 5 minutes. Note: Upgrade operations times are O(minutes)", pollingDeadlineFlag)
	}
	if o.pollingInterval > o.pollingDeadline {
		return fmt.Errorf("--%s=%v must be greater than --%s=%v", pollingDeadlineFlag, o.pollingDeadline, pollingIntervalFlag, o.pollingInterval)
	}

	// Version validation.
	if (o.desiredControlPlaneVersion == "" && !o.inPlaceControlPlaneUpgrade) ||
		(o.desiredControlPlaneVersion != "" && o.inPlaceControlPlaneUpgrade) {
		return fmt.Errorf("specify --%s or provide a version for --%s, but not both", inPlaceControlPlaneUpgradeFlag, desiredControlPlaneVersionFlag)
	}
	if o.desiredControlPlaneVersion != "" {
		if err := clusters.IsFormatValid(o.desiredControlPlaneVersion); err != nil {
			return fmt.Errorf("--%s=%q is not valid: %v", desiredControlPlaneVersionFlag, o.desiredControlPlaneVersion, err)
		}
	}
	if err := clusters.IsFormatValid(o.desiredNodeVersion); err != nil {
		return fmt.Errorf("--%s=%q is not valid: %v", desiredNodeVersionFlag, o.desiredNodeVersion, err)
	}
	// Use of `-` or `latest` aliases are validated later at the control plane and node pool level.
	if !o.inPlaceControlPlaneUpgrade &&
		o.desiredControlPlaneVersion != clusters.DefaultVersion &&
		o.desiredControlPlaneVersion != clusters.LatestVersion &&
		o.desiredNodeVersion != clusters.DefaultVersion &&
		o.desiredNodeVersion != clusters.LatestVersion {
		if err := clusters.IsWithinVersionSkew(o.desiredNodeVersion, o.desiredControlPlaneVersion, clusters.MaxVersionSkew); err != nil {
			return err
		}
	}

	return nil
}

// Complete cascades down the resource hierarchy, ensuring that all descendent migrators are initialized.
func (o *migrateOptions) Complete(ctx context.Context) error {
	authedClient, err := google.DefaultClient(ctx, compute.CloudPlatformScope)
	if err != nil {
		return err
	}

	o.clients, err = o.fetchClientFunc(ctx, o.containerBasePath, authedClient)
	if err != nil {
		return err
	}

	handler := operations.NewHandler(o.pollingInterval, o.pollingDeadline)
	options := &clusters.Options{
		ConcurrentNodePools:        ConcurrentNodePools,
		DesiredControlPlaneVersion: o.desiredControlPlaneVersion,
		DesiredNodeVersion:         o.desiredNodeVersion,
		InPlaceControlPlaneUpgrade: o.inPlaceControlPlaneUpgrade,
	}

	factory := func(n *compute.Network) migrate.Migrator {
		return networks.New(o.projectID, n, handler, o.clients, o.concurrentClusters, options)
	}

	log.Infof("Fetching network %s for project %q", o.selectedNetwork, o.projectID)

	ns, err := o.clients.Compute.ListNetworks(ctx, o.projectID)
	if err != nil {
		return fmt.Errorf("error listing networks: %w", err)
	}

	o.migrators = make([]migrate.Migrator, 0)
	for _, n := range ns {
		if n.Name == o.selectedNetwork {
			o.migrators = append(o.migrators, factory(n))
		}
	}

	if len(o.migrators) == 0 {
		return fmt.Errorf("unable to find network %s", o.selectedNetwork)
	}

	return nil
}

// Run cascades down the resource hierarchy, initializing, validating, and converting all descendent migrators.
func (o *migrateOptions) Run(ctx context.Context) error {
	sem := make(chan struct{}, ConcurrentNetworks)

	log.Info("Initialize objects for conversion.")
	if err := migrate.Complete(ctx, sem, o.migrators...); err != nil {
		return err
	}

	log.Info("Validate resources for conversion.")
	if err := migrate.Validate(ctx, sem, o.migrators...); err != nil {
		return err
	}

	if o.validateOnly {
		log.Infof("--%s==true; skipping conversion.", validateOnlyFlag)
		return nil
	}

	log.Info("Initiate resource conversion.")
	return migrate.Migrate(ctx, sem, o.migrators...)
}

// setupCloseHandler cancels the shared context when the user hits ctrl-c.
func setupCloseHandler(cancel context.CancelFunc) {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-c
		cancel()
	}()
}

func getRetryableClientOption(retry int, waitMin, waitMax time.Duration, authedClient *http.Client) option.ClientOption {
	retryClient := retryablehttp.NewClient()
	retryClient.RetryMax = retry
	retryClient.RetryWaitMin = waitMin
	retryClient.RetryWaitMax = waitMax
	retryClient.Logger = nil
	retryClient.CheckRetry = retryPolicy()

	c := retryClient.StandardClient()
	c.Transport.(*retryablehttp.RoundTripper).Client.HTTPClient = authedClient
	return option.WithHTTPClient(c)
}

func fetchClients(ctx context.Context, basePath string, authedClient *http.Client) (*pkg.Clients, error) {
	opt := getRetryableClientOption(3, 5*time.Second, 30*time.Second, authedClient)
	computeService, err := compute.NewService(ctx, opt)
	if err != nil {
		return nil, err
	}
	containerService, err := container.NewService(ctx, opt)
	if err != nil {
		return nil, err
	}

	// Retry for up-to 5 minutes for Compute Alpha API calls.
	alphaOpt := getRetryableClientOption(5, 5*time.Second, 160*time.Second, authedClient)
	computeServiceAlpha, err := computealpha.NewService(ctx, alphaOpt)
	if err != nil {
		return nil, err
	}

	if basePath != "" {
		containerService.BasePath = basePath
	}
	return &pkg.Clients{
		Compute: &pkg.Compute{
			V1:    computeService,
			Alpha: computeServiceAlpha,
		},
		Container: &pkg.Container{V1: containerService},
	}, nil
}

func retryPolicy() func(ctx context.Context, resp *http.Response, err error) (bool, error) {
	return func(ctx context.Context, resp *http.Response, err error) (bool, error) {
		shouldRetry, newErr := retryablehttp.DefaultRetryPolicy(ctx, resp, err)
		if newErr != nil || shouldRetry == true {
			return shouldRetry, newErr
		}

		// GCE returns 403 as a RateLimiting response code.
		if resp.StatusCode == 403 {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return true, err
			}
			return true, errors.New(string(body))
		}
		return false, nil
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"
	"os"

	"github.com/dtan4/k8ship/kubernetes"
	"github.com/spf13/cobra"
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	Use:   "k8ship",
	Short: "Ship image to Kubernetes easily",
}

var rootOpts = struct {
	context    string
	kubeconfig string
}{}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	RootCmd.PersistentFlags().StringVar(&rootOpts.context, "context", "", "Kubernetes context")
	RootCmd.PersistentFlags().StringVar(&rootOpts.kubeconfig, "kubeconfig", "", "kubeconfig path")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if rootOpts.kubeconfig == "" {
		if os.Getenv("KUBECONFIG") == "" {
			rootOpts.kubeconfig = kubernetes.DefaultConfigFile()
		} else {
			rootOpts.kubeconfig = os.Getenv("KUBECONFIG")
		}
	}
}
<commit_msg>Suppress error output<commit_after>package cmd

import (
	"fmt"
	"os"

	"github.com/dtan4/k8ship/kubernetes"
	"github.com/spf13/cobra"
)

// RootCmd represents the base command when called without any subcommands
var RootCmd = &cobra.Command{
	SilenceErrors: true,
	SilenceUsage:  true,
	Use:           "k8ship",
	Short:         "Ship image to Kubernetes easily",
}

var rootOpts = struct {
	context    string
	kubeconfig string
}{}

// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
	if err := RootCmd.Execute(); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)

	RootCmd.PersistentFlags().StringVar(&rootOpts.context, "context", "", "Kubernetes context")
	RootCmd.PersistentFlags().StringVar(&rootOpts.kubeconfig, "kubeconfig", "", "kubeconfig path")
}

// initConfig reads in config file and ENV variables if set.
func initConfig() {
	if rootOpts.kubeconfig == "" {
		if os.Getenv("KUBECONFIG") == "" {
			rootOpts.kubeconfig = kubernetes.DefaultConfigFile()
		} else {
			rootOpts.kubeconfig = os.Getenv("KUBECONFIG")
		}
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"./apis"
	"./apps"
	"./developers"
	"./env"
	"./org"
	"./products"
	"./shared"
	"./sharedflows"
	"./sync"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var RootCmd = &cobra.Command{
	Use:     "apigeeapi",
	Version: "0.1",
	Short:   "Utility to work with Apigee APIs.",
	Long:    "This command lets you interact with Apigee APIs.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {

		shared.Init()

		if shared.RootArgs.Token == "" && shared.RootArgs.ServiceAccount == "" {
			return fmt.Errorf("Either token or service account must be provided")
		} else {
			if shared.RootArgs.ServiceAccount != "" {
				viper.SetConfigFile(shared.RootArgs.ServiceAccount)
				err := viper.ReadInConfig() // Find and read the config file
				if err != nil {             // Handle errors reading the config file
					return fmt.Errorf("Fatal error config file: %s \n", err)
				} else {
					if viper.Get("private_key") == "" {
						return fmt.Errorf("Fatal error: Private key missing in the service account")
					}
					if viper.Get("client_email") == "" {
						return fmt.Errorf("Fatal error: client email missing in the service account")
					}
					_, err = shared.GenerateAccessToken()
					if err != nil {
						return fmt.Errorf("Fatal error generating access token: %s \n", err)
					} else {
						return nil
					}
				}
			}
			return nil
		}
	},
}

func init() {
	cobra.OnInitialize(initConfig)

	RootCmd.PersistentFlags().BoolVarP(&shared.LogInfo, "log", "l",
		false, "Log Information")

	RootCmd.PersistentFlags().StringVarP(&shared.RootArgs.Token, "token", "t",
		"", "Google OAuth Token")
	viper.BindPFlag("token", RootCmd.PersistentFlags().Lookup("token"))

	RootCmd.PersistentFlags().StringVarP(&shared.RootArgs.ServiceAccount, "account", "a",
		"", "Path Service Account private key in JSON")
	viper.BindPFlag("account", RootCmd.PersistentFlags().Lookup("account"))

	RootCmd.AddCommand(apis.Cmd)
	RootCmd.AddCommand(org.Cmd)
	RootCmd.AddCommand(sync.Cmd)
	RootCmd.AddCommand(env.Cmd)
	RootCmd.AddCommand(products.Cmd)
	RootCmd.AddCommand(developers.Cmd)
	RootCmd.AddCommand(apps.Cmd)
	RootCmd.AddCommand(sharedflows.Cmd)
}

func initConfig() {
	viper.AutomaticEnv() // read in environment variables that match
	viper.SetConfigType("json")

}

// GetRootCmd returns the root of the cobra command-tree.
func GetRootCmd() *cobra.Command {
	return RootCmd
}
<commit_msg>set env prefix<commit_after>package cmd

import (
	"./apis"
	"./apps"
	"./developers"
	"./env"
	"./org"
	"./products"
	"./shared"
	"./sharedflows"
	"./sync"
	"fmt"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var RootCmd = &cobra.Command{
	Use:     "apigeeapi",
	Version: "0.1",
	Short:   "Utility to work with Apigee APIs.",
	Long:    "This command lets you interact with Apigee APIs.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error {

		shared.Init()

		if shared.RootArgs.Token == "" && shared.RootArgs.ServiceAccount == "" {
			return fmt.Errorf("Either token or service account must be provided")
		} else {
			if shared.RootArgs.ServiceAccount != "" {
				viper.SetConfigFile(shared.RootArgs.ServiceAccount)
				err := viper.ReadInConfig() // Find and read the config file
				if err != nil {             // Handle errors reading the config file
					return fmt.Errorf("Fatal error config file: %s \n", err)
				} else {
					if viper.Get("private_key") == "" {
						return fmt.Errorf("Fatal error: Private key missing in the service account")
					}
					if viper.Get("client_email") == "" {
						return fmt.Errorf("Fatal error: client email missing in the service account")
					}
					_, err = shared.GenerateAccessToken()
					if err != nil {
						return fmt.Errorf("Fatal error generating access token: %s \n", err)
					} else {
						return nil
					}
				}
			}
			return nil
		}
	},
}

func init() {
	cobra.OnInitialize(initConfig)

	RootCmd.PersistentFlags().BoolVarP(&shared.LogInfo, "log", "l",
		false, "Log Information")

	RootCmd.PersistentFlags().StringVarP(&shared.RootArgs.Token, "token", "t",
		"", "Google OAuth Token")
	viper.BindPFlag("token", RootCmd.PersistentFlags().Lookup("token"))

	RootCmd.PersistentFlags().StringVarP(&shared.RootArgs.ServiceAccount, "account", "a",
		"", "Path Service Account private key in JSON")
	viper.BindPFlag("account", RootCmd.PersistentFlags().Lookup("account"))

	RootCmd.AddCommand(apis.Cmd)
	RootCmd.AddCommand(org.Cmd)
	RootCmd.AddCommand(sync.Cmd)
	RootCmd.AddCommand(env.Cmd)
	RootCmd.AddCommand(products.Cmd)
	RootCmd.AddCommand(developers.Cmd)
	RootCmd.AddCommand(apps.Cmd)
	RootCmd.AddCommand(sharedflows.Cmd)
}

func initConfig() {
	viper.SetEnvPrefix("APIGEE")
	viper.AutomaticEnv() // read in environment variables that match
	viper.SetConfigType("json")
}

// GetRootCmd returns the root of the cobra command-tree.
func GetRootCmd() *cobra.Command {
	return RootCmd
}
<|endoftext|>


<commit_before>// Copyright © 2017 Karl Hepworth <Karl.Hepworth@gmail.com>
//
// 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 cmd

import (
	"log"
	"os"
	"os/exec"

	"fmt"

	"github.com/spf13/cobra"
)

// syncCmd represents the backup command
var syncCmd = &cobra.Command{
	Use:   "sync",
	Short: "Execute drush sql-sync or rsync between two drush aliases",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		d, err := exec.LookPath("drush")
		if err != nil {
			log.Fatal("Drush was not found in your $PATH")
		}
		var yesVal string
		if yes {
			yesVal = "--yes";
		}
		if syncFiles {
			{
				fsPu := fmt.Sprintf("@%v:%%files", source)
				fdPu := fmt.Sprintf("@%v:%%files", destination)
				c := exec.Command(d, yesVal, "rsync", fsPu, fdPu, "--exclude-other-sites", "--exclude-conf")
				c.Stdin = os.Stdin
				c.Stdout = os.Stdout
				c.Stderr = os.Stderr
				c.Run()
				c.Wait()
			}
			{
				fsPr := fmt.Sprintf("@%v:%%private", source)
				fdPr := fmt.Sprintf("@%v:%%private", destination)
				c := exec.Command(d, yesVal, "rsync", fsPr, fdPr, "--exclude-other-sites", "--exclude-conf")
				c.Stdin = os.Stdin
				c.Stdout = os.Stdout
				c.Stderr = os.Stderr
				c.Run()
				c.Wait()
			}
		}
		if syncDatabase {
			c := exec.Command(d, "sql-sync", source, destination, "--yes")
			c.Stdin = os.Stdin
			c.Stdout = os.Stdout
			c.Stderr = os.Stderr
			c.Run()
			c.Wait()
		}
	},
}

func init() {
	RootCmd.AddCommand(syncCmd)

	// Get the current working directory.
	dir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}

	syncCmd.Flags().StringVarP(&source, "source", "s", "", "Drush alias to use as source")
	syncCmd.Flags().StringVarP(&destination, "destination", "d", dir, "Drush alias to use as destination")
	syncCmd.Flags().BoolVarP(&syncDatabase, "database", "b", false, "Flag database for sync action.")
	syncCmd.Flags().BoolVarP(&syncFiles, "files", "f", false, "Flag files for sync action.")
	syncCmd.Flags().BoolVarP(&yes, "yes", "y", false, "Use command with --yes")

	syncCmd.MarkFlagRequired("source")
	syncCmd.MarkFlagRequired("destination")
}
<commit_msg>Improve the usefulness of the sync command, and the applicability of the --yes flag.<commit_after>// Copyright © 2017 Karl Hepworth <Karl.Hepworth@gmail.com>
//
// 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 cmd

import (
	"log"
	"os"
	"os/exec"

	"fmt"

	"github.com/spf13/cobra"
)

// syncCmd represents the backup command
var syncCmd = &cobra.Command{
	Use:   "sync",
	Short: "Execute drush sql-sync or rsync between two drush aliases",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {
		d, err := exec.LookPath("drush")
		if err != nil {
			log.Fatal("Drush was not found in your $PATH")
		}
		if syncFiles {
			{
				if name == "" {
					log.Fatal("rsync target was not provided, please specify target with --target.")
				}
				fsPu := fmt.Sprintf("%v:%%%v", source, name)
				fdPu := fmt.Sprintf("%v:%%%v", destination, name)
				if (yes) {
					c := exec.Command(d, "rsync", fsPu, fdPu, "--exclude-other-sites", "--exclude-conf", "--yes")
					c.Stdin = os.Stdin
					c.Stdout = os.Stdout
					c.Stderr = os.Stderr
					c.Run()
					c.Wait()
				} else {
					c := exec.Command(d, "rsync", fsPu, fdPu, "--exclude-other-sites", "--exclude-conf")
					c.Stdin = os.Stdin
					c.Stdout = os.Stdout
					c.Stderr = os.Stderr
					c.Run()
					c.Wait()
				}
			}
		}
		if syncDatabase {
			if yes {
				c := exec.Command(d, "sql-sync", source, destination, "--yes")
				c.Stdin = os.Stdin
				c.Stdout = os.Stdout
				c.Stderr = os.Stderr
				c.Run()
				c.Wait()
			} else {
				c := exec.Command(d, "sql-sync", source, destination)
				c.Stdin = os.Stdin
				c.Stdout = os.Stdout
				c.Stderr = os.Stderr
				c.Run()
				c.Wait()
			}
		}
	},
}

func init() {
	RootCmd.AddCommand(syncCmd)

	// Get the current working directory.
	dir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}

	syncCmd.Flags().StringVarP(&source, "source", "s", "", "Drush alias to use as source")
	syncCmd.Flags().StringVarP(&destination, "destination", "d", dir, "Drush alias to use as destination")
	syncCmd.Flags().StringVarP(&name, "target", "t", "", "The name of the path alias in the drush alias. ie files, public, private, temp")
	syncCmd.Flags().BoolVarP(&syncDatabase, "database", "b", false, "Flag database for sync action.")
	syncCmd.Flags().BoolVarP(&syncFiles, "files", "f", false, "Flag files for sync action.")
	syncCmd.Flags().BoolVarP(&yes, "yes", "y", false, "Use command with --yes")

	syncCmd.MarkFlagRequired("source")
	syncCmd.MarkFlagRequired("destination")
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageProjectRegex = "[a-z](?:[-a-z0-9]{0,61}[a-z0-9])?" // TODO(paddy): this isn't based on any documentation; we're just copying the image name restrictions. Need to follow up with @danawillow and/or @evandbrown and see if there's an actual limit to this
	resolveImageFamilyRegex  = "[a-z](?:[-a-z0-9]{0,61}[a-z0-9])?" // TODO(paddy): this isn't based on any documentation; we're just copying the image name restrictions. Need to follow up with @danawillow and/or @evandbrown and see if there's an actual limit to this
	resolveImageImageRegex   = "[a-z](?:[-a-z0-9]{0,61}[a-z0-9])?" // 1-63 characters, lowercase letters, numbers, and hyphens only, beginning and ending in a lowercase letter or number
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("^projects/(%s)/global/images/(%s)$", resolveImageProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("^projects/(%s)/global/images/family/(%s)$", resolveImageProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", resolveImageProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", resolveImageProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/v1/projects/(%s)/global/images/(%s)", resolveImageProjectRegex, resolveImageImageRegex))
)

func resolveImageImageExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, name string) (string, error) {
	// built-in projects to look for images/families containing the string
	// on the left in
	imageMap := map[string]string{
		"centos":   "centos-cloud",
		"coreos":   "coreos-cloud",
		"debian":   "debian-cloud",
		"opensuse": "opensuse-cloud",
		"rhel":     "rhel-cloud",
		"sles":     "suse-cloud",
		"ubuntu":   "ubuntu-os-cloud",
		"windows":  "windows-cloud",
	}
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if len(res)-1 != 2 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d project image regex matches, got %d for %s", 2, len(res)-1, name)
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if len(res)-1 != 2 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d project family regex matches, got %d for %s", 2, len(res)-1, name)
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if len(res)-1 != 1 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d global image regex matches, got %d for %s", 1, len(res)-1, name)
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if len(res)-1 != 1 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d global family regex matches, got %d for %s", 1, len(res)-1, name)
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if len(res)-1 != 1 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d family family regex matches, got %d for %s", 1, len(res)-1, name)
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if len(res)-1 != 2 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d project image shorthand regex matches, got %d for %s", 2, len(res)-1, name)
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if len(res)-1 != 2 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d project family shorthand regex matches, got %d for %s", 2, len(res)-1, name)
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if len(res)-1 != 1 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d image regex matches, got %d for %s", 1, len(res)-1, name)
		}
		if ok, err := resolveImageImageExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if len(res)-1 != 1 { // subtract one, index zero is the entire matched expression
			return "", fmt.Errorf("Expected %d family regex matches, got %d for %s", 1, len(res)-1, name)
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}
<commit_msg>Update with @danawillow's feedback.<commit_after>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageProjectRegex = "[-_a-zA-Z0-9]*"
	resolveImageFamilyRegex  = "[-_a-zA-Z0-9]*"
	resolveImageImageRegex   = "[-_a-zA-Z0-9]*"
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("^projects/(%s)/global/images/(%s)$", resolveImageProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("^projects/(%s)/global/images/family/(%s)$", resolveImageProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", resolveImageProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", resolveImageProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/v1/projects/(%s)/global/images/(%s)", resolveImageProjectRegex, resolveImageImageRegex))
)

func resolveImageImageExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

func sanityTestRegexMatches(expected int, got []string, regexType, name string) error {
	if len(got)-1 != expected { // subtract one, index zero is the entire matched expression
		return fmt.Errorf("Expected %d %s regex matches, got %d for %s", 2, regexType, len(got)-1, name)
	}
	return nil
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, name string) (string, error) {
	// built-in projects to look for images/families containing the string
	// on the left in
	imageMap := map[string]string{
		"centos":   "centos-cloud",
		"coreos":   "coreos-cloud",
		"debian":   "debian-cloud",
		"opensuse": "opensuse-cloud",
		"rhel":     "rhel-cloud",
		"sles":     "suse-cloud",
		"ubuntu":   "ubuntu-os-cloud",
		"windows":  "windows-cloud",
	}
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "image", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}
<|endoftext|>


<commit_before>package blackjack

import "log"

// ParseCard returns the integer value of a card following blackjack ruleset.
func ParseCard(card string) int {
	switch card {
	case "one":
		return 1
	case "two":
		return 2
	case "three":
		return 3
	case "four":
		return 4
	case "five":
		return 5
	case "six":
		return 6
	case "seven":
		return 7
	case "eight":
		return 8
	case "nine":
		return 9
	case "ten", "jack", "queen", "king":
		return 10
	case "ace":
		return 11
	default:
		log.Printf("card %s is not a valid card", card)
		return 0
	}
}

// IsBlackjack returns true if the player has a blackjack, false otherwise.
func IsBlackjack(card1, card2 string) bool {
	panic("Please implement the IsBlackjack function")
}

// LargeHand implements the decision tree for hand scores larger than 20 points.
func LargeHand(isBlackjack bool, dealerScore int) string {
	panic("Please implement the LargeHand function")
}

// SmallHand implements the decision tree for hand scores with less than 21 points.
func SmallHand(handScore, dealerScore int) string {
	panic("Please implement the SmallHand function")
}

// FirstTurn returns the semi-optimal decision for the first turn, given the cards of the player and the dealer.
// This function is already implemented and does not need to be edited. It pulls the other functions together in a
// complete decision tree for the first turn.
func FirstTurn(card1, card2, dealerCard string) string {
	handScore := ParseCard(card1) + ParseCard(card2)
	dealerScore := ParseCard(dealerCard)

	if 20 < handScore {
		return LargeHand(IsBlackjack(card1, card2), dealerScore)
	}
	return SmallHand(handScore, dealerScore)
}
<commit_msg>Implmenet IsBlackjack<commit_after>package blackjack

import "log"

// ParseCard returns the integer value of a card following blackjack ruleset.
func ParseCard(card string) int {
	switch card {
	case "one":
		return 1
	case "two":
		return 2
	case "three":
		return 3
	case "four":
		return 4
	case "five":
		return 5
	case "six":
		return 6
	case "seven":
		return 7
	case "eight":
		return 8
	case "nine":
		return 9
	case "ten", "jack", "queen", "king":
		return 10
	case "ace":
		return 11
	default:
		log.Printf("card %s is not a valid card", card)
		return 0
	}
}

// IsBlackjack returns true if the player has a blackjack, false otherwise.
func IsBlackjack(card1, card2 string) bool {
	return ParseCard(card1)+ParseCard(card2) == 21
}

// LargeHand implements the decision tree for hand scores larger than 20 points.
func LargeHand(isBlackjack bool, dealerScore int) string {
	panic("Please implement the LargeHand function")
}

// SmallHand implements the decision tree for hand scores with less than 21 points.
func SmallHand(handScore, dealerScore int) string {
	panic("Please implement the SmallHand function")
}

// FirstTurn returns the semi-optimal decision for the first turn, given the cards of the player and the dealer.
// This function is already implemented and does not need to be edited. It pulls the other functions together in a
// complete decision tree for the first turn.
func FirstTurn(card1, card2, dealerCard string) string {
	handScore := ParseCard(card1) + ParseCard(card2)
	dealerScore := ParseCard(dealerCard)

	if 20 < handScore {
		return LargeHand(IsBlackjack(card1, card2), dealerScore)
	}
	return SmallHand(handScore, dealerScore)
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package client

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"

	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/exec/exectest"
	"github.com/tsuru/tsuru/fs/fstest"
	"gopkg.in/check.v1"
)

func (s *S) TestPluginInstallInfo(c *check.C) {
	c.Assert(PluginInstall{}.Info(), check.NotNil)
}

func (s *S) TestPluginInstall(c *check.C) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "fakeplugin")
	}))
	defer ts.Close()
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", ts.URL},
		Stdout: &stdout,
	}
	client := cmd.NewClient(nil, nil, manager)
	command := PluginInstall{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	pluginsPath := cmd.JoinWithUserDir(".tsuru", "plugins")
	hasAction := rfs.HasAction(fmt.Sprintf("mkdirall %s with mode 0755", pluginsPath))
	c.Assert(hasAction, check.Equals, true)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	hasAction = rfs.HasAction(fmt.Sprintf("openfile %s with mode 0755", pluginPath))
	c.Assert(hasAction, check.Equals, true)
	f, err := rfs.Open(pluginPath)
	c.Assert(err, check.IsNil)
	data, err := ioutil.ReadAll(f)
	c.Assert(err, check.IsNil)
	c.Assert("fakeplugin\n", check.Equals, string(data))
	expected := `Plugin "myplugin" successfully installed!` + "\n"
	c.Assert(expected, check.Equals, stdout.String())
}

func (s *S) TestPluginInstallError(c *check.C) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("my err"))
	}))
	defer ts.Close()
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", ts.URL},
		Stdout: &stdout,
	}
	client := cmd.NewClient(nil, nil, manager)
	command := PluginInstall{}
	err := command.Run(&context, client)
	c.Assert(err, check.ErrorMatches, `Invalid status code reading plugin: 500 - "my err"`)
}

func (s *S) TestPluginInstallIsACommand(c *check.C) {
	var _ cmd.Command = &PluginInstall{}
}

func (s *S) TestPlugin(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"a", "b"}), check.Equals, true)
	c.Assert(buf.String(), check.Equals, "hello world")
	commands := fexec.GetCommands(pluginPath)
	c.Assert(commands, check.HasLen, 1)
	target, err := cmd.GetTarget()
	c.Assert(err, check.IsNil)
	token, err := cmd.ReadToken()
	c.Assert(err, check.IsNil)
	envs := os.Environ()
	tsuruEnvs := []string{
		fmt.Sprintf("TSURU_TARGET=%s", target),
		fmt.Sprintf("TSURU_TOKEN=%s", token),
		"TSURU_PLUGIN_NAME=myplugin",
	}
	envs = append(envs, tsuruEnvs...)
	c.Assert(commands[0].GetEnvs(), check.DeepEquals, envs)
}

func (s *S) TestPluginWithArgs(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	context := cmd.Context{Args: []string{"myplugin", "ble", "bla"}}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"ble", "bla"}), check.Equals, true)
}

func (s *S) TestPluginTryNameWithAnyExtension(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"otherplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "otherplugin.exe")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"a", "b"}), check.Equals, true)
	c.Assert(buf.String(), check.Equals, "hello world")
	commands := fexec.GetCommands(pluginPath)
	c.Assert(commands, check.HasLen, 1)
	target, err := cmd.GetTarget()
	c.Assert(err, check.IsNil)
	token, err := cmd.ReadToken()
	c.Assert(err, check.IsNil)
	envs := os.Environ()
	tsuruEnvs := []string{
		fmt.Sprintf("TSURU_TARGET=%s", target),
		fmt.Sprintf("TSURU_TOKEN=%s", token),
		"TSURU_PLUGIN_NAME=otherplugin",
	}
	envs = append(envs, tsuruEnvs...)
	c.Assert(commands[0].GetEnvs(), check.DeepEquals, envs)
}

func (s *S) TestPluginLoop(c *check.C) {
	os.Setenv("TSURU_PLUGIN_NAME", "myplugin")
	defer os.Unsetenv("TSURU_PLUGIN_NAME")
	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.Equals, cmd.ErrLookup)
}

func (s *S) TestPluginCommandNotFound(c *check.C) {
	fexec := exectest.ErrorExecutor{Err: os.ErrNotExist}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.Equals, cmd.ErrLookup)
}

func (s *S) TestPluginRemoveInfo(c *check.C) {
	c.Assert(PluginRemove{}.Info(), check.NotNil)
}

func (s *S) TestPluginRemove(c *check.C) {
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin"},
		Stdout: &stdout,
	}
	client := cmd.NewClient(nil, nil, manager)
	command := PluginRemove{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	hasAction := rfs.HasAction(fmt.Sprintf("remove %s", pluginPath))
	c.Assert(hasAction, check.Equals, true)
	expected := `Plugin "myplugin" successfully removed!` + "\n"
	c.Assert(expected, check.Equals, stdout.String())
}

func (s *S) TestPluginRemoveIsACommand(c *check.C) {
	var _ cmd.Command = &PluginRemove{}
}

func (s *S) TestPluginListInfo(c *check.C) {
	c.Assert(PluginList{}.Info(), check.NotNil)
}

func (s *S) TestPluginListIsACommand(c *check.C) {
	var _ cmd.Command = &PluginList{}
}
<commit_msg>tsuru: use empty http.Client in plugin tests<commit_after>// Copyright 2016 tsuru-client authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package client

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"

	"github.com/tsuru/tsuru/cmd"
	"github.com/tsuru/tsuru/exec/exectest"
	"github.com/tsuru/tsuru/fs/fstest"
	"gopkg.in/check.v1"
)

func (s *S) TestPluginInstallInfo(c *check.C) {
	c.Assert(PluginInstall{}.Info(), check.NotNil)
}

func (s *S) TestPluginInstall(c *check.C) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "fakeplugin")
	}))
	defer ts.Close()
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", ts.URL},
		Stdout: &stdout,
	}
	client := cmd.NewClient(&http.Client{}, nil, manager)
	command := PluginInstall{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	pluginsPath := cmd.JoinWithUserDir(".tsuru", "plugins")
	hasAction := rfs.HasAction(fmt.Sprintf("mkdirall %s with mode 0755", pluginsPath))
	c.Assert(hasAction, check.Equals, true)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	hasAction = rfs.HasAction(fmt.Sprintf("openfile %s with mode 0755", pluginPath))
	c.Assert(hasAction, check.Equals, true)
	f, err := rfs.Open(pluginPath)
	c.Assert(err, check.IsNil)
	data, err := ioutil.ReadAll(f)
	c.Assert(err, check.IsNil)
	c.Assert("fakeplugin\n", check.Equals, string(data))
	expected := `Plugin "myplugin" successfully installed!` + "\n"
	c.Assert(expected, check.Equals, stdout.String())
}

func (s *S) TestPluginInstallError(c *check.C) {
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte("my err"))
	}))
	defer ts.Close()
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", ts.URL},
		Stdout: &stdout,
	}
	client := cmd.NewClient(&http.Client{}, nil, manager)
	command := PluginInstall{}
	err := command.Run(&context, client)
	c.Assert(err, check.ErrorMatches, `Invalid status code reading plugin: 500 - "my err"`)
}

func (s *S) TestPluginInstallIsACommand(c *check.C) {
	var _ cmd.Command = &PluginInstall{}
}

func (s *S) TestPlugin(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"a", "b"}), check.Equals, true)
	c.Assert(buf.String(), check.Equals, "hello world")
	commands := fexec.GetCommands(pluginPath)
	c.Assert(commands, check.HasLen, 1)
	target, err := cmd.GetTarget()
	c.Assert(err, check.IsNil)
	token, err := cmd.ReadToken()
	c.Assert(err, check.IsNil)
	envs := os.Environ()
	tsuruEnvs := []string{
		fmt.Sprintf("TSURU_TARGET=%s", target),
		fmt.Sprintf("TSURU_TOKEN=%s", token),
		"TSURU_PLUGIN_NAME=myplugin",
	}
	envs = append(envs, tsuruEnvs...)
	c.Assert(commands[0].GetEnvs(), check.DeepEquals, envs)
}

func (s *S) TestPluginWithArgs(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	context := cmd.Context{Args: []string{"myplugin", "ble", "bla"}}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"ble", "bla"}), check.Equals, true)
}

func (s *S) TestPluginTryNameWithAnyExtension(c *check.C) {
	// Kids, do not try this at $HOME
	defer os.Setenv("HOME", os.Getenv("HOME"))
	tempHome, _ := filepath.Abs("testdata")
	os.Setenv("HOME", tempHome)

	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"otherplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "otherplugin.exe")
	c.Assert(fexec.ExecutedCmd(pluginPath, []string{"a", "b"}), check.Equals, true)
	c.Assert(buf.String(), check.Equals, "hello world")
	commands := fexec.GetCommands(pluginPath)
	c.Assert(commands, check.HasLen, 1)
	target, err := cmd.GetTarget()
	c.Assert(err, check.IsNil)
	token, err := cmd.ReadToken()
	c.Assert(err, check.IsNil)
	envs := os.Environ()
	tsuruEnvs := []string{
		fmt.Sprintf("TSURU_TARGET=%s", target),
		fmt.Sprintf("TSURU_TOKEN=%s", token),
		"TSURU_PLUGIN_NAME=otherplugin",
	}
	envs = append(envs, tsuruEnvs...)
	c.Assert(commands[0].GetEnvs(), check.DeepEquals, envs)
}

func (s *S) TestPluginLoop(c *check.C) {
	os.Setenv("TSURU_PLUGIN_NAME", "myplugin")
	defer os.Unsetenv("TSURU_PLUGIN_NAME")
	fexec := exectest.FakeExecutor{
		Output: map[string][][]byte{
			"a b": {[]byte("hello world")},
		},
	}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.Equals, cmd.ErrLookup)
}

func (s *S) TestPluginCommandNotFound(c *check.C) {
	fexec := exectest.ErrorExecutor{Err: os.ErrNotExist}
	Execut = &fexec
	defer func() {
		Execut = nil
	}()
	var buf bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin", "a", "b"},
		Stdout: &buf,
		Stderr: &buf,
	}
	err := RunPlugin(&context)
	c.Assert(err, check.Equals, cmd.ErrLookup)
}

func (s *S) TestPluginRemoveInfo(c *check.C) {
	c.Assert(PluginRemove{}.Info(), check.NotNil)
}

func (s *S) TestPluginRemove(c *check.C) {
	rfs := fstest.RecordingFs{}
	fsystem = &rfs
	defer func() {
		fsystem = nil
	}()
	var stdout bytes.Buffer
	context := cmd.Context{
		Args:   []string{"myplugin"},
		Stdout: &stdout,
	}
	client := cmd.NewClient(&http.Client{}, nil, manager)
	command := PluginRemove{}
	err := command.Run(&context, client)
	c.Assert(err, check.IsNil)
	pluginPath := cmd.JoinWithUserDir(".tsuru", "plugins", "myplugin")
	hasAction := rfs.HasAction(fmt.Sprintf("remove %s", pluginPath))
	c.Assert(hasAction, check.Equals, true)
	expected := `Plugin "myplugin" successfully removed!` + "\n"
	c.Assert(expected, check.Equals, stdout.String())
}

func (s *S) TestPluginRemoveIsACommand(c *check.C) {
	var _ cmd.Command = &PluginRemove{}
}

func (s *S) TestPluginListInfo(c *check.C) {
	c.Assert(PluginList{}.Info(), check.NotNil)
}

func (s *S) TestPluginListIsACommand(c *check.C) {
	var _ cmd.Command = &PluginList{}
}
<|endoftext|>


<commit_before><commit_msg>Check in<commit_after><|endoftext|>


<commit_before>package libkb

import (
	"testing"
	"time"

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

const chWait = 5 * time.Second

func TestBgTicker(t *testing.T) {
	duration := 2 * time.Millisecond
	start := time.Now()
	ticker := NewBgTickerWithWait(duration, time.Millisecond)

	// Test tick
	for i := 0; i < 5; i++ {
		select {
		case <-ticker.C:
			if i == 0 {
				require.True(t, time.Since(start) >= duration)
			}
		case <-time.After(chWait):
			require.Fail(t, "ticker did not fire")
		}
	}
}
<commit_msg>Fix BgTickerTest<commit_after>package libkb

import (
	"testing"
	"time"

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

const chWait = 5 * time.Second

func TestBgTicker(t *testing.T) {
	duration := 2 * time.Millisecond
	wait := time.Millisecond
	start := time.Now()
	ticker := NewBgTickerWithWait(duration, wait)

	// Test tick
	for i := 0; i < 5; i++ {
		select {
		case <-ticker.C:
			if i == 0 {
				require.True(t, time.Since(start) >= wait, "time.Since(start) %v", time.Since(start))
			}
		case <-time.After(chWait):
			require.Fail(t, "ticker did not fire")
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"sort"
	"sync"

	"code.google.com/p/go.net/context"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/topo"
)

type versionDebugVars struct {
	BuildHost      string
	BuildUser      string
	BuildTimestamp int64
	BuildGitRev    string
}

func (wr *Wrangler) GetVersion(tabletAlias topo.TabletAlias) (string, error) {
	tablet, err := wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return "", err
	}

	resp, err := http.Get("http://" + tablet.Addr() + "/debug/vars")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	vars := versionDebugVars{}
	err = json.Unmarshal(body, &vars)
	if err != nil {
		return "", err
	}

	version := fmt.Sprintf("%v", vars)

	log.Infof("Tablet %v is running version '%v'", tabletAlias, version)
	return version, nil
}

// helper method to asynchronously get and diff a version
func (wr *Wrangler) diffVersion(masterVersion string, masterAlias topo.TabletAlias, alias topo.TabletAlias, wg *sync.WaitGroup, er concurrency.ErrorRecorder) {
	defer wg.Done()
	log.Infof("Gathering version for %v", alias)
	slaveVersion, err := wr.GetVersion(alias)
	if err != nil {
		er.RecordError(err)
		return
	}

	if masterVersion != slaveVersion {
		er.RecordError(fmt.Errorf("Master %v version %v is different than slave %v version %v", masterAlias, masterVersion, alias, slaveVersion))
	}
}

func (wr *Wrangler) ValidateVersionShard(keyspace, shard string) error {
	si, err := wr.ts.GetShard(keyspace, shard)
	if err != nil {
		return err
	}

	// get version from the master, or error
	if si.MasterAlias.Uid == topo.NO_TABLET {
		return fmt.Errorf("No master in shard %v/%v", keyspace, shard)
	}
	log.Infof("Gathering version for master %v", si.MasterAlias)
	masterVersion, err := wr.GetVersion(si.MasterAlias)
	if err != nil {
		return err
	}

	// read all the aliases in the shard, that is all tablets that are
	// replicating from the master
	aliases, err := topo.FindAllTabletAliasesInShard(context.TODO(), wr.ts, keyspace, shard)
	if err != nil {
		return err
	}

	// then diff with all slaves
	er := concurrency.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, alias := range aliases {
		if alias == si.MasterAlias {
			continue
		}

		wg.Add(1)
		go wr.diffVersion(masterVersion, si.MasterAlias, alias, &wg, &er)
	}
	wg.Wait()
	if er.HasErrors() {
		return fmt.Errorf("Version diffs:\n%v", er.Error().Error())
	}
	return nil
}

func (wr *Wrangler) ValidateVersionKeyspace(keyspace string) error {
	// find all the shards
	shards, err := wr.ts.GetShardNames(keyspace)
	if err != nil {
		return err
	}

	// corner cases
	if len(shards) == 0 {
		return fmt.Errorf("No shards in keyspace %v", keyspace)
	}
	sort.Strings(shards)
	if len(shards) == 1 {
		return wr.ValidateVersionShard(keyspace, shards[0])
	}

	// find the reference version using the first shard's master
	si, err := wr.ts.GetShard(keyspace, shards[0])
	if err != nil {
		return err
	}
	if si.MasterAlias.Uid == topo.NO_TABLET {
		return fmt.Errorf("No master in shard %v/%v", keyspace, shards[0])
	}
	referenceAlias := si.MasterAlias
	log.Infof("Gathering version for reference master %v", referenceAlias)
	referenceVersion, err := wr.GetVersion(referenceAlias)
	if err != nil {
		return err
	}

	// then diff with all tablets but master 0
	er := concurrency.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, shard := range shards {
		aliases, err := topo.FindAllTabletAliasesInShard(context.TODO(), wr.ts, keyspace, shard)
		if err != nil {
			er.RecordError(err)
			continue
		}

		for _, alias := range aliases {
			if alias == si.MasterAlias {
				continue
			}

			wg.Add(1)
			go wr.diffVersion(referenceVersion, referenceAlias, alias, &wg, &er)
		}
	}
	wg.Wait()
	if er.HasErrors() {
		return fmt.Errorf("Version diffs:\n%v", er.Error().Error())
	}
	return nil
}
<commit_msg>Make wrangler version determination pluggable.<commit_after>// Copyright 2013, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package wrangler

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"sort"
	"sync"

	"code.google.com/p/go.net/context"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/concurrency"
	"github.com/youtube/vitess/go/vt/topo"
)

var getVersionFromTablet = func(tabletAddr string) (string, error) {
	resp, err := http.Get("http://" + tabletAddr + "/debug/vars")
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	var vars struct {
		BuildHost      string
		BuildUser      string
		BuildTimestamp int64
		BuildGitRev    string
	}
	err = json.Unmarshal(body, &vars)
	if err != nil {
		return "", err
	}

	version := fmt.Sprintf("%v", vars)
	return version, nil
}

func (wr *Wrangler) GetVersion(tabletAlias topo.TabletAlias) (string, error) {
	tablet, err := wr.ts.GetTablet(tabletAlias)
	if err != nil {
		return "", err
	}

	version, err := getVersionFromTablet(tablet.Addr())
	if err != nil {
		return "", err
	}
	log.Infof("Tablet %v is running version '%v'", tabletAlias, version)
	return version, err
}

// helper method to asynchronously get and diff a version
func (wr *Wrangler) diffVersion(masterVersion string, masterAlias topo.TabletAlias, alias topo.TabletAlias, wg *sync.WaitGroup, er concurrency.ErrorRecorder) {
	defer wg.Done()
	log.Infof("Gathering version for %v", alias)
	slaveVersion, err := wr.GetVersion(alias)
	if err != nil {
		er.RecordError(err)
		return
	}

	if masterVersion != slaveVersion {
		er.RecordError(fmt.Errorf("Master %v version %v is different than slave %v version %v", masterAlias, masterVersion, alias, slaveVersion))
	}
}

func (wr *Wrangler) ValidateVersionShard(keyspace, shard string) error {
	si, err := wr.ts.GetShard(keyspace, shard)
	if err != nil {
		return err
	}

	// get version from the master, or error
	if si.MasterAlias.Uid == topo.NO_TABLET {
		return fmt.Errorf("No master in shard %v/%v", keyspace, shard)
	}
	log.Infof("Gathering version for master %v", si.MasterAlias)
	masterVersion, err := wr.GetVersion(si.MasterAlias)
	if err != nil {
		return err
	}

	// read all the aliases in the shard, that is all tablets that are
	// replicating from the master
	aliases, err := topo.FindAllTabletAliasesInShard(context.TODO(), wr.ts, keyspace, shard)
	if err != nil {
		return err
	}

	// then diff with all slaves
	er := concurrency.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, alias := range aliases {
		if alias == si.MasterAlias {
			continue
		}

		wg.Add(1)
		go wr.diffVersion(masterVersion, si.MasterAlias, alias, &wg, &er)
	}
	wg.Wait()
	if er.HasErrors() {
		return fmt.Errorf("Version diffs:\n%v", er.Error().Error())
	}
	return nil
}

func (wr *Wrangler) ValidateVersionKeyspace(keyspace string) error {
	// find all the shards
	shards, err := wr.ts.GetShardNames(keyspace)
	if err != nil {
		return err
	}

	// corner cases
	if len(shards) == 0 {
		return fmt.Errorf("No shards in keyspace %v", keyspace)
	}
	sort.Strings(shards)
	if len(shards) == 1 {
		return wr.ValidateVersionShard(keyspace, shards[0])
	}

	// find the reference version using the first shard's master
	si, err := wr.ts.GetShard(keyspace, shards[0])
	if err != nil {
		return err
	}
	if si.MasterAlias.Uid == topo.NO_TABLET {
		return fmt.Errorf("No master in shard %v/%v", keyspace, shards[0])
	}
	referenceAlias := si.MasterAlias
	log.Infof("Gathering version for reference master %v", referenceAlias)
	referenceVersion, err := wr.GetVersion(referenceAlias)
	if err != nil {
		return err
	}

	// then diff with all tablets but master 0
	er := concurrency.AllErrorRecorder{}
	wg := sync.WaitGroup{}
	for _, shard := range shards {
		aliases, err := topo.FindAllTabletAliasesInShard(context.TODO(), wr.ts, keyspace, shard)
		if err != nil {
			er.RecordError(err)
			continue
		}

		for _, alias := range aliases {
			if alias == si.MasterAlias {
				continue
			}

			wg.Add(1)
			go wr.diffVersion(referenceVersion, referenceAlias, alias, &wg, &er)
		}
	}
	wg.Wait()
	if er.HasErrors() {
		return fmt.Errorf("Version diffs:\n%v", er.Error().Error())
	}
	return nil
}
<|endoftext|>


<commit_before>package triplestore

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/d4l3k/messagediff"
	"github.com/degdb/degdb/protocol"
)

func TestTripleStore(t *testing.T) {
	file, err := ioutil.TempFile(os.TempDir(), "triplestore.db")
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(file.Name())
	if err := initDB(file.Name()); err != nil {
		t.Fatal(err)
	}

	triples := []*protocol.Triple{
		{
			Subj: "/m/02mjmr",
			Pred: "/type/object/name",
			Obj:  "Barack Obama",
		},
		{
			Subj: "/m/02mjmr",
			Pred: "/type/object/type",
			Obj:  "/people/person",
		},
		{
			Subj: "/m/0hume",
			Pred: "/type/object/name",
			Obj:  "Hume",
		},
		{
			Subj: "/m/0hume",
			Pred: "/type/object/type",
			Obj:  "/organization/team",
		},
	}

	Insert(triples)
	// Insert twice to ensure no duplicates.
	Insert(triples)

	testData := []struct {
		query *protocol.Triple
		want  []*protocol.Triple
	}{
		{
			&protocol.Triple{
				Subj: "/m/02mjmr",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/name",
					Obj:  "Barack Obama",
				},
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/type",
					Obj:  "/people/person",
				},
			},
		},
		{
			&protocol.Triple{
				Pred: "/type/object/type",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/type",
					Obj:  "/people/person",
				},
				{
					Subj: "/m/0hume",
					Pred: "/type/object/type",
					Obj:  "/organization/team",
				},
			},
		},
		{
			&protocol.Triple{
				Pred: "/type/object/name",
				Obj:  "Barack Obama",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/name",
					Obj:  "Barack Obama",
				},
			},
		},
	}

	for i, td := range testData {
		triples, err := Query(td.query, -1)
		if err != nil {
			t.Error(err)
		}
		if diff, ok := messagediff.PrettyDiff(td.want, triples); !ok {
			t.Errorf("%d. Query(%#v, -1) = %#v; diff %s", i, td.query, triples, diff)
		}
	}

	info, err := Size()
	if err != nil {
		t.Fatal(err)
	}
	if info.Triples != uint64(len(triples)) {
		t.Errorf("Size() = %#v; not %d", info, len(triples))
	}
}
<commit_msg>Fixed triplestore tests after refactor<commit_after>package triplestore

import (
	"io/ioutil"
	"os"
	"testing"

	"github.com/d4l3k/messagediff"
	"github.com/degdb/degdb/protocol"
)

func TestTripleStore(t *testing.T) {
	file, err := ioutil.TempFile(os.TempDir(), "triplestore.db")
	if err != nil {
		t.Fatal(err)
	}
	defer os.Remove(file.Name())
	db, err := NewTripleStore(file.Name())
	if err != nil {
		t.Fatal(err)
	}

	triples := []*protocol.Triple{
		{
			Subj: "/m/02mjmr",
			Pred: "/type/object/name",
			Obj:  "Barack Obama",
		},
		{
			Subj: "/m/02mjmr",
			Pred: "/type/object/type",
			Obj:  "/people/person",
		},
		{
			Subj: "/m/0hume",
			Pred: "/type/object/name",
			Obj:  "Hume",
		},
		{
			Subj: "/m/0hume",
			Pred: "/type/object/type",
			Obj:  "/organization/team",
		},
	}

	db.Insert(triples)
	// Insert twice to ensure no duplicates.
	db.Insert(triples)

	testData := []struct {
		query *protocol.Triple
		want  []*protocol.Triple
	}{
		{
			&protocol.Triple{
				Subj: "/m/02mjmr",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/name",
					Obj:  "Barack Obama",
				},
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/type",
					Obj:  "/people/person",
				},
			},
		},
		{
			&protocol.Triple{
				Pred: "/type/object/type",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/type",
					Obj:  "/people/person",
				},
				{
					Subj: "/m/0hume",
					Pred: "/type/object/type",
					Obj:  "/organization/team",
				},
			},
		},
		{
			&protocol.Triple{
				Pred: "/type/object/name",
				Obj:  "Barack Obama",
			},
			[]*protocol.Triple{
				{
					Subj: "/m/02mjmr",
					Pred: "/type/object/name",
					Obj:  "Barack Obama",
				},
			},
		},
	}

	for i, td := range testData {
		triples, err := db.Query(td.query, -1)
		if err != nil {
			t.Error(err)
		}
		if diff, ok := messagediff.PrettyDiff(td.want, triples); !ok {
			t.Errorf("%d. Query(%#v, -1) = %#v; diff %s", i, td.query, triples, diff)
		}
	}

	info, err := db.Size()
	if err != nil {
		t.Fatal(err)
	}
	if info.Triples != uint64(len(triples)) {
		t.Errorf("Size() = %#v; not %d", info, len(triples))
	}
}
<|endoftext|>


<commit_before>// Copyright 2015, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

package flate

import (
	"io"

	"github.com/dsnet/compress/internal"
	"github.com/dsnet/compress/internal/prefix"
)

type Reader struct {
	InputOffset  int64 // Total number of bytes read from underlying io.Reader
	OutputOffset int64 // Total number of bytes emitted from Read

	rd     prefixReader // Input source
	toRead []byte       // Uncompressed data ready to be emitted from Read
	dist   int          // The current distance
	blkLen int          // Uncompressed bytes left to read in meta-block
	cpyLen int          // Bytes left to backward dictionary copy
	last   bool         // Last block bit detected
	err    error        // Persistent error

	step      func(*Reader) // Single step of decompression work (can panic)
	stepState int           // The sub-step state for certain steps

	dict     dictDecoder    // Dynamic sliding dictionary
	litTree  prefix.Decoder // Literal and length symbol prefix decoder
	distTree prefix.Decoder // Backward distance symbol prefix decoder
}

func NewReader(r io.Reader) *Reader {
	zr := new(Reader)
	zr.Reset(r)
	return zr
}

func (zr *Reader) Read(buf []byte) (int, error) {
	for {
		if len(zr.toRead) > 0 {
			cnt := copy(buf, zr.toRead)
			zr.toRead = zr.toRead[cnt:]
			zr.OutputOffset += int64(cnt)
			return cnt, nil
		}
		if zr.err != nil {
			return 0, zr.err
		}

		// Perform next step in decompression process.
		zr.rd.Offset = zr.InputOffset
		func() {
			defer errRecover(&zr.err)
			zr.step(zr)
		}()
		var err error
		if zr.InputOffset, err = zr.rd.Flush(); err != nil {
			zr.err = err
		}
		if zr.err != nil {
			if zr.err == internal.ErrInvalid {
				zr.err = ErrCorrupt
			}
		}
		if zr.err != nil {
			zr.toRead = zr.dict.ReadFlush() // Flush what's left in case of error
		}
	}
}

func (zr *Reader) Close() error {
	if zr.err == io.EOF || zr.err == ErrClosed {
		zr.toRead = nil // Make sure future reads fail
		zr.err = ErrClosed
		return nil
	}
	return zr.err // Return the persistent error
}

func (zr *Reader) Reset(r io.Reader) {
	*zr = Reader{
		rd:   zr.rd,
		step: (*Reader).readBlockHeader,
		dict: zr.dict,
	}
	zr.rd.Init(r)
	zr.dict.Init(maxHistSize)
}

// readBlockHeader reads the block header according to RFC section 3.2.3.
func (zr *Reader) readBlockHeader() {
	if zr.last {
		zr.rd.ReadPads()
		panic(io.EOF)
	}

	zr.last = zr.rd.ReadBits(1) == 1
	switch zr.rd.ReadBits(2) {
	case 0:
		// Raw block (RFC section 3.2.4).
		zr.rd.ReadPads()

		n := uint16(zr.rd.ReadBits(16))
		nn := uint16(zr.rd.ReadBits(16))
		if n^nn != 0xffff {
			panic(ErrCorrupt)
		}
		zr.blkLen = int(n)

		// By convention, an empty block flushes the read buffer.
		if zr.blkLen == 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlockHeader
			return
		}
		zr.step = (*Reader).readRawData
	case 1:
		// Fixed prefix block (RFC section 3.2.6).
		zr.litTree, zr.distTree = decLit, decDist
		zr.step = (*Reader).readBlock
	case 2:
		// Dynamic prefix block (RFC section 3.2.7).
		zr.rd.ReadPrefixCodes(&zr.litTree, &zr.distTree)
		zr.step = (*Reader).readBlock
	default:
		// Reserved block (RFC section 3.2.3).
		panic(ErrCorrupt)
	}
}

// readRawData reads raw data according to RFC section 3.2.4.
func (zr *Reader) readRawData() {
	buf := zr.dict.WriteSlice()
	if len(buf) > zr.blkLen {
		buf = buf[:zr.blkLen]
	}

	cnt, err := zr.rd.Read(buf)
	zr.blkLen -= cnt
	zr.dict.WriteMark(cnt)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		panic(err)
	}

	if zr.blkLen > 0 {
		zr.toRead = zr.dict.ReadFlush()
		zr.step = (*Reader).readRawData // We need to continue this work
		return
	}
	zr.step = (*Reader).readBlockHeader
}

// readCommands reads block commands according to RFC section 3.2.3.
func (zr *Reader) readBlock() {
	const (
		stateInit = iota // Zero value must be stateInit
		stateDict
	)

	switch zr.stepState {
	case stateInit:
		goto readLiteral
	case stateDict:
		goto copyDistance
	}

readLiteral:
	// Read literal and/or (length, distance) according to RFC section 3.2.3.
	{
		if zr.dict.AvailSize() == 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlock
			zr.stepState = stateInit // Need to continue work here
			return
		}

		// Read the literal symbol.
		litSym, ok := zr.rd.TryReadSymbol(&zr.litTree)
		if !ok {
			litSym = zr.rd.ReadSymbol(&zr.litTree)
		}
		switch {
		case litSym < endBlockSym:
			zr.dict.WriteByte(byte(litSym))
			goto readLiteral
		case litSym == endBlockSym:
			zr.step = (*Reader).readBlockHeader
			zr.stepState = stateInit // Next call to readBlock must start here
			return
		case litSym < maxNumLitSyms:
			// Decode the copy length.
			rec := lenRanges[litSym-257]
			extra, ok := zr.rd.TryReadBits(uint(rec.Len))
			if !ok {
				extra = zr.rd.ReadBits(uint(rec.Len))
			}
			zr.cpyLen = int(rec.Base) + int(extra)

			// Read the distance symbol.
			distSym, ok := zr.rd.TryReadSymbol(&zr.distTree)
			if !ok {
				distSym = zr.rd.ReadSymbol(&zr.distTree)
			}
			if distSym >= maxNumDistSyms {
				panic(ErrCorrupt)
			}

			// Decode the copy distance.
			rec = distRanges[distSym]
			extra, ok = zr.rd.TryReadBits(uint(rec.Len))
			if !ok {
				extra = zr.rd.ReadBits(uint(rec.Len))
			}
			zr.dist = int(rec.Base) + int(extra)

			goto copyDistance
		default:
			panic(ErrCorrupt)
		}
	}

copyDistance:
	// Perform a backwards copy according to RFC section 3.2.3.
	{
		cnt := zr.dict.TryWriteCopy(zr.dist, zr.cpyLen)
		if cnt == 0 {
			cnt = zr.dict.WriteCopy(zr.dist, zr.cpyLen)
		}
		zr.cpyLen -= cnt

		if zr.cpyLen > 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlock
			zr.stepState = stateDict // Need to continue work here
			return
		} else {
			goto readLiteral
		}
	}
}
<commit_msg>flate: move Reader.Reset closer to NewReader<commit_after>// Copyright 2015, Joe Tsai. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE.md file.

package flate

import (
	"io"

	"github.com/dsnet/compress/internal"
	"github.com/dsnet/compress/internal/prefix"
)

type Reader struct {
	InputOffset  int64 // Total number of bytes read from underlying io.Reader
	OutputOffset int64 // Total number of bytes emitted from Read

	rd     prefixReader // Input source
	toRead []byte       // Uncompressed data ready to be emitted from Read
	dist   int          // The current distance
	blkLen int          // Uncompressed bytes left to read in meta-block
	cpyLen int          // Bytes left to backward dictionary copy
	last   bool         // Last block bit detected
	err    error        // Persistent error

	step      func(*Reader) // Single step of decompression work (can panic)
	stepState int           // The sub-step state for certain steps

	dict     dictDecoder    // Dynamic sliding dictionary
	litTree  prefix.Decoder // Literal and length symbol prefix decoder
	distTree prefix.Decoder // Backward distance symbol prefix decoder
}

func NewReader(r io.Reader) *Reader {
	zr := new(Reader)
	zr.Reset(r)
	return zr
}

func (zr *Reader) Reset(r io.Reader) {
	*zr = Reader{
		rd:   zr.rd,
		step: (*Reader).readBlockHeader,
		dict: zr.dict,
	}
	zr.rd.Init(r)
	zr.dict.Init(maxHistSize)
}

func (zr *Reader) Read(buf []byte) (int, error) {
	for {
		if len(zr.toRead) > 0 {
			cnt := copy(buf, zr.toRead)
			zr.toRead = zr.toRead[cnt:]
			zr.OutputOffset += int64(cnt)
			return cnt, nil
		}
		if zr.err != nil {
			return 0, zr.err
		}

		// Perform next step in decompression process.
		zr.rd.Offset = zr.InputOffset
		func() {
			defer errRecover(&zr.err)
			zr.step(zr)
		}()
		var err error
		if zr.InputOffset, err = zr.rd.Flush(); err != nil {
			zr.err = err
		}
		if zr.err != nil {
			if zr.err == internal.ErrInvalid {
				zr.err = ErrCorrupt
			}
		}
		if zr.err != nil {
			zr.toRead = zr.dict.ReadFlush() // Flush what's left in case of error
		}
	}
}

func (zr *Reader) Close() error {
	if zr.err == io.EOF || zr.err == ErrClosed {
		zr.toRead = nil // Make sure future reads fail
		zr.err = ErrClosed
		return nil
	}
	return zr.err // Return the persistent error
}

// readBlockHeader reads the block header according to RFC section 3.2.3.
func (zr *Reader) readBlockHeader() {
	if zr.last {
		zr.rd.ReadPads()
		panic(io.EOF)
	}

	zr.last = zr.rd.ReadBits(1) == 1
	switch zr.rd.ReadBits(2) {
	case 0:
		// Raw block (RFC section 3.2.4).
		zr.rd.ReadPads()

		n := uint16(zr.rd.ReadBits(16))
		nn := uint16(zr.rd.ReadBits(16))
		if n^nn != 0xffff {
			panic(ErrCorrupt)
		}
		zr.blkLen = int(n)

		// By convention, an empty block flushes the read buffer.
		if zr.blkLen == 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlockHeader
			return
		}
		zr.step = (*Reader).readRawData
	case 1:
		// Fixed prefix block (RFC section 3.2.6).
		zr.litTree, zr.distTree = decLit, decDist
		zr.step = (*Reader).readBlock
	case 2:
		// Dynamic prefix block (RFC section 3.2.7).
		zr.rd.ReadPrefixCodes(&zr.litTree, &zr.distTree)
		zr.step = (*Reader).readBlock
	default:
		// Reserved block (RFC section 3.2.3).
		panic(ErrCorrupt)
	}
}

// readRawData reads raw data according to RFC section 3.2.4.
func (zr *Reader) readRawData() {
	buf := zr.dict.WriteSlice()
	if len(buf) > zr.blkLen {
		buf = buf[:zr.blkLen]
	}

	cnt, err := zr.rd.Read(buf)
	zr.blkLen -= cnt
	zr.dict.WriteMark(cnt)
	if err != nil {
		if err == io.EOF {
			err = io.ErrUnexpectedEOF
		}
		panic(err)
	}

	if zr.blkLen > 0 {
		zr.toRead = zr.dict.ReadFlush()
		zr.step = (*Reader).readRawData // We need to continue this work
		return
	}
	zr.step = (*Reader).readBlockHeader
}

// readCommands reads block commands according to RFC section 3.2.3.
func (zr *Reader) readBlock() {
	const (
		stateInit = iota // Zero value must be stateInit
		stateDict
	)

	switch zr.stepState {
	case stateInit:
		goto readLiteral
	case stateDict:
		goto copyDistance
	}

readLiteral:
	// Read literal and/or (length, distance) according to RFC section 3.2.3.
	{
		if zr.dict.AvailSize() == 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlock
			zr.stepState = stateInit // Need to continue work here
			return
		}

		// Read the literal symbol.
		litSym, ok := zr.rd.TryReadSymbol(&zr.litTree)
		if !ok {
			litSym = zr.rd.ReadSymbol(&zr.litTree)
		}
		switch {
		case litSym < endBlockSym:
			zr.dict.WriteByte(byte(litSym))
			goto readLiteral
		case litSym == endBlockSym:
			zr.step = (*Reader).readBlockHeader
			zr.stepState = stateInit // Next call to readBlock must start here
			return
		case litSym < maxNumLitSyms:
			// Decode the copy length.
			rec := lenRanges[litSym-257]
			extra, ok := zr.rd.TryReadBits(uint(rec.Len))
			if !ok {
				extra = zr.rd.ReadBits(uint(rec.Len))
			}
			zr.cpyLen = int(rec.Base) + int(extra)

			// Read the distance symbol.
			distSym, ok := zr.rd.TryReadSymbol(&zr.distTree)
			if !ok {
				distSym = zr.rd.ReadSymbol(&zr.distTree)
			}
			if distSym >= maxNumDistSyms {
				panic(ErrCorrupt)
			}

			// Decode the copy distance.
			rec = distRanges[distSym]
			extra, ok = zr.rd.TryReadBits(uint(rec.Len))
			if !ok {
				extra = zr.rd.ReadBits(uint(rec.Len))
			}
			zr.dist = int(rec.Base) + int(extra)

			goto copyDistance
		default:
			panic(ErrCorrupt)
		}
	}

copyDistance:
	// Perform a backwards copy according to RFC section 3.2.3.
	{
		cnt := zr.dict.TryWriteCopy(zr.dist, zr.cpyLen)
		if cnt == 0 {
			cnt = zr.dict.WriteCopy(zr.dist, zr.cpyLen)
		}
		zr.cpyLen -= cnt

		if zr.cpyLen > 0 {
			zr.toRead = zr.dict.ReadFlush()
			zr.step = (*Reader).readBlock
			zr.stepState = stateDict // Need to continue work here
			return
		} else {
			goto readLiteral
		}
	}
}
<|endoftext|>