<commit_before>// Copyright 2018 Google Inc. All rights reserved.

// 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 (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/GoogleCloudPlatform/container-structure-test/pkg/drivers"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/output"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/types"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/types/unversioned"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/utils"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
)

const warnMessage = `WARNING: Using the host driver runs tests directly on the machine
that this binary is being run on, and can potentially corrupt your system.
Be sure you know what you're doing before continuing!

Continue? (y/n)`

var totalTests int

var driverImpl func(drivers.DriverConfig) (drivers.Driver, error)
var args *drivers.DriverConfig

var TestCmd = &cobra.Command{
	Use:   "test",
	Short: "Runs the tests",
	Long:  `Runs the tests`,
	Args: func(cmd *cobra.Command, _ []string) error {
		return validateArgs()
	},
	Run: func(cmd *cobra.Command, _ []string) {
		testResults := Run()
		out := &output.OutWriter{
			Verbose: verbose,
			Quiet:   quiet,
		}
		os.Exit(out.OutputResults(testResults))
	},
}

func validateArgs() error {
	if driver == drivers.Host {
		if metadata == "" {
			return fmt.Errorf("Please provide path to image metadata file")
		}
		if imagePath != "" {
			return fmt.Errorf("Cannot provide both image path and metadata file")
		}
	} else {
		if imagePath == "" {
			return fmt.Errorf("Please supply path to image or tarball to test against")
		}
		if metadata != "" {
			return fmt.Errorf("Cannot provide both image path and metadata file")
		}
	}
	if len(configFiles) == 0 {
		return fmt.Errorf("Please provide at least one test config file")
	}
	return nil
}

func RunTests() []*unversioned.FullResult {
	fullResults := []*unversioned.FullResult{}
	for _, file := range configFiles {
		tests, err := Parse(file)
		if err != nil {
			logrus.Fatalf("Error parsing config file: %s", err)
		}
		fullResults = append(fullResults, &unversioned.FullResult{
			FileName: file,
			Results:  tests.RunAll(),
		})
	}
	return fullResults
}

func Parse(fp string) (types.StructureTest, error) {
	testContents, err := ioutil.ReadFile(fp)
	if err != nil {
		return nil, err
	}

	// We first have to unmarshal to determine the schema version, then we unmarshal again
	// to do the full parse.
	var unmarshal types.Unmarshaller
	var strictUnmarshal types.Unmarshaller
	var versionHolder types.SchemaVersion

	switch {
	case strings.HasSuffix(fp, ".json"):
		unmarshal = json.Unmarshal
		strictUnmarshal = json.Unmarshal
	case strings.HasSuffix(fp, ".yaml"):
		unmarshal = yaml.Unmarshal
		strictUnmarshal = yaml.UnmarshalStrict
	default:
		return nil, errors.New("Please provide valid JSON or YAML config file")
	}

	if err := unmarshal(testContents, &versionHolder); err != nil {
		return nil, err
	}

	version := versionHolder.SchemaVersion
	if version == "" {
		return nil, errors.New("Please provide JSON schema version")
	}

	var st types.StructureTest
	if schemaVersion, ok := types.SchemaVersions[version]; ok {
		st = schemaVersion()
	} else {
		return nil, errors.New("Unsupported schema version: " + version)
	}

	strictUnmarshal(testContents, st)

	tests, _ := st.(types.StructureTest) //type assertion
	tests.SetDriverImpl(driverImpl, *args)
	return tests, nil
}

func Run() []*unversioned.FullResult {
	args = &drivers.DriverConfig{
		Image:    imagePath,
		Save:     save,
		Metadata: metadata,
	}

	var err error

	if pull {
		if driver != drivers.Docker {
			logrus.Fatal("Image pull not supported when not using Docker driver")
		}
		var repository, tag string
		parts := strings.Split(imagePath, ":")
		if len(parts) < 2 {
			logrus.Fatal("Please provide specific tag for image")
		}
		repository = parts[0]
		tag = parts[1]
		client, err := docker.NewClientFromEnv()
		if err = client.PullImage(docker.PullImageOptions{
			Repository:   repository,
			Tag:          tag,
			OutputStream: os.Stdout,
		}, docker.AuthConfiguration{}); err != nil {
			logrus.Fatalf("Error pulling remote image %s: %s", imagePath, err.Error())
		}
	}

	if driver == drivers.Host && !utils.UserConfirmation(warnMessage, force) {
		os.Exit(1)
	}

	driverImpl = drivers.InitDriverImpl(driver)
	if driverImpl == nil {
		logrus.Fatalf("Unsupported driver type: %s", driver)
	}
	if err != nil {
		logrus.Fatal(err.Error())
	}
	logrus.Infof("Using driver %s\n", driver)
	return RunTests()
}

func init() {
	RootCmd.AddCommand(TestCmd)
	TestCmd.Flags().StringVar(&imagePath, "image", "", "path to test image")
	TestCmd.Flags().StringVar(&driver, "driver", "docker", "driver to use when running tests")
	TestCmd.Flags().StringVar(&metadata, "metadata", "", "path to image metadata file")
	TestCmd.Flags().BoolVar(&pull, "pull", false, "force a pull of the image before running tests")
	TestCmd.Flags().BoolVar(&save, "save", false, "preserve created containers after test run")
	TestCmd.Flags().BoolVar(&force, "force", false, "force run of host driver (without command line input)")
	TestCmd.Flags().BoolVar(&quiet, "quiet", false, "flag to suppress output")
	TestCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "verbose testing output")
	TestCmd.Flags().StringArrayVar(&configFiles, "config", []string{}, "test config files")
}
<commit_msg>allow .yml file extension (#91)<commit_after>// Copyright 2018 Google Inc. All rights reserved.

// 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 (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	"github.com/GoogleCloudPlatform/container-structure-test/pkg/drivers"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/output"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/types"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/types/unversioned"
	"github.com/GoogleCloudPlatform/container-structure-test/pkg/utils"
	docker "github.com/fsouza/go-dockerclient"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
)

const warnMessage = `WARNING: Using the host driver runs tests directly on the machine
that this binary is being run on, and can potentially corrupt your system.
Be sure you know what you're doing before continuing!

Continue? (y/n)`

var totalTests int

var driverImpl func(drivers.DriverConfig) (drivers.Driver, error)
var args *drivers.DriverConfig

var TestCmd = &cobra.Command{
	Use:   "test",
	Short: "Runs the tests",
	Long:  `Runs the tests`,
	Args: func(cmd *cobra.Command, _ []string) error {
		return validateArgs()
	},
	Run: func(cmd *cobra.Command, _ []string) {
		testResults := Run()
		out := &output.OutWriter{
			Verbose: verbose,
			Quiet:   quiet,
		}
		os.Exit(out.OutputResults(testResults))
	},
}

func validateArgs() error {
	if driver == drivers.Host {
		if metadata == "" {
			return fmt.Errorf("Please provide path to image metadata file")
		}
		if imagePath != "" {
			return fmt.Errorf("Cannot provide both image path and metadata file")
		}
	} else {
		if imagePath == "" {
			return fmt.Errorf("Please supply path to image or tarball to test against")
		}
		if metadata != "" {
			return fmt.Errorf("Cannot provide both image path and metadata file")
		}
	}
	if len(configFiles) == 0 {
		return fmt.Errorf("Please provide at least one test config file")
	}
	return nil
}

func RunTests() []*unversioned.FullResult {
	fullResults := []*unversioned.FullResult{}
	for _, file := range configFiles {
		tests, err := Parse(file)
		if err != nil {
			logrus.Fatalf("Error parsing config file: %s", err)
		}
		fullResults = append(fullResults, &unversioned.FullResult{
			FileName: file,
			Results:  tests.RunAll(),
		})
	}
	return fullResults
}

func Parse(fp string) (types.StructureTest, error) {
	testContents, err := ioutil.ReadFile(fp)
	if err != nil {
		return nil, err
	}

	// We first have to unmarshal to determine the schema version, then we unmarshal again
	// to do the full parse.
	var unmarshal types.Unmarshaller
	var strictUnmarshal types.Unmarshaller
	var versionHolder types.SchemaVersion

	switch {
	case strings.HasSuffix(fp, ".json"):
		unmarshal = json.Unmarshal
		strictUnmarshal = json.Unmarshal
	case strings.HasSuffix(fp, ".yaml"):
		unmarshal = yaml.Unmarshal
		strictUnmarshal = yaml.UnmarshalStrict
	case strings.HasSuffix(fp, ".yml"):
		unmarshal = yaml.Unmarshal
		strictUnmarshal = yaml.UnmarshalStrict
	default:
		return nil, errors.New("Please provide valid JSON or YAML config file")
	}

	if err := unmarshal(testContents, &versionHolder); err != nil {
		return nil, err
	}

	version := versionHolder.SchemaVersion
	if version == "" {
		return nil, errors.New("Please provide JSON schema version")
	}

	var st types.StructureTest
	if schemaVersion, ok := types.SchemaVersions[version]; ok {
		st = schemaVersion()
	} else {
		return nil, errors.New("Unsupported schema version: " + version)
	}

	strictUnmarshal(testContents, st)

	tests, _ := st.(types.StructureTest) //type assertion
	tests.SetDriverImpl(driverImpl, *args)
	return tests, nil
}

func Run() []*unversioned.FullResult {
	args = &drivers.DriverConfig{
		Image:    imagePath,
		Save:     save,
		Metadata: metadata,
	}

	var err error

	if pull {
		if driver != drivers.Docker {
			logrus.Fatal("Image pull not supported when not using Docker driver")
		}
		var repository, tag string
		parts := strings.Split(imagePath, ":")
		if len(parts) < 2 {
			logrus.Fatal("Please provide specific tag for image")
		}
		repository = parts[0]
		tag = parts[1]
		client, err := docker.NewClientFromEnv()
		if err = client.PullImage(docker.PullImageOptions{
			Repository:   repository,
			Tag:          tag,
			OutputStream: os.Stdout,
		}, docker.AuthConfiguration{}); err != nil {
			logrus.Fatalf("Error pulling remote image %s: %s", imagePath, err.Error())
		}
	}

	if driver == drivers.Host && !utils.UserConfirmation(warnMessage, force) {
		os.Exit(1)
	}

	driverImpl = drivers.InitDriverImpl(driver)
	if driverImpl == nil {
		logrus.Fatalf("Unsupported driver type: %s", driver)
	}
	if err != nil {
		logrus.Fatal(err.Error())
	}
	logrus.Infof("Using driver %s\n", driver)
	return RunTests()
}

func init() {
	RootCmd.AddCommand(TestCmd)
	TestCmd.Flags().StringVar(&imagePath, "image", "", "path to test image")
	TestCmd.Flags().StringVar(&driver, "driver", "docker", "driver to use when running tests")
	TestCmd.Flags().StringVar(&metadata, "metadata", "", "path to image metadata file")
	TestCmd.Flags().BoolVar(&pull, "pull", false, "force a pull of the image before running tests")
	TestCmd.Flags().BoolVar(&save, "save", false, "preserve created containers after test run")
	TestCmd.Flags().BoolVar(&force, "force", false, "force run of host driver (without command line input)")
	TestCmd.Flags().BoolVar(&quiet, "quiet", false, "flag to suppress output")
	TestCmd.Flags().BoolVarP(&verbose, "verbose", "v", false, "verbose testing output")
	TestCmd.Flags().StringArrayVar(&configFiles, "config", []string{}, "test config files")
}
<|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"
	"net/mail"

	"os"
	"strings"

	"fmt"

	alias2 "github.com/fubarhouse/drubuild/alias"
	"github.com/fubarhouse/drubuild/user"
	"github.com/spf13/cobra"
)

// userCmd represents the user command
var userCmd = &cobra.Command{
	Use:   "user",
	Short: "User management tasks",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {

		if !strings.Contains(pattern, "%v") {
			log.Fatalf("Specified pattern '%v' does not include alias modifier '%%v'.", pattern)
			os.Exit(1)
		}

		Email, err := mail.ParseAddress(user_email)
		if err != nil {
			log.Fatal("Input email address could not be parsed.")
			log.Fatal(err)
		}

		if user_block {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					if User.State == 1 {
						User.State = 0
						User.StateChange()
					} else {
						log.Printf("User '%v' is already blocked on %v\n", User.Name, Alias)
					}
				} else {
					log.Printf("User '%v' was not found on %v\n", User.Name, Alias)
				}
			}
		}

		if user_create {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				User.Alias = Alias
				User.Name = user_name
				User.Email = user_email
				User.Create(user_password)
				if user_role != "" {
					if !User.HasRole(user_role) {
						User.Roles = append(User.Roles, user_role)
					}
					User.RolesAdd()
				}
			}
		}
		if user_delete {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					User.Delete()
				} else {
					log.Printf("User '%v' was not found on %v", user_name, Alias)
				}
			}
		}

		if user_unblock {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					if User.State == 0 {
						User.State = 1
						User.StateChange()
					} else {
						log.Printf("User '%v' is already unblocked on %v", User.Name, Alias)
					}
				} else {
					log.Printf("User '%v' was not found on %v", user_name, Alias)
				}
			}
		}

		if user_verify {
			for _, ThisAlias := range strings.Split(aliases, ",") {
				ThisAlias = strings.Trim(ThisAlias, " ")
				ThisAlias = strings.Replace(pattern, "%v", ThisAlias, 1)
				log.Printf("Beginning to work with alias %v.", ThisAlias)
				UserGroup := user.NewDrupalUserGroup()
				Alias := alias2.NewAlias(ThisAlias, "", ThisAlias)
				if Alias.GetStatus() {
					UserGroup.Populate(Alias.GetName())
					User := UserGroup.GetUser(user_name)
					if User.Alias == "" {
						User.Alias = Alias.GetName()
					}
					if User.Name == "" {
						User.Name = user_name
					}
					User.Email = Email.Address

					if user_create {
						User.Create(user_password)
					} else {
						log.Println("User creation not set, skipping.")
					}

					if user_block {
						User.State = 1
						User.StateChange()
					} else if user_unblock {
						User.State = 0
						User.StateChange()
					} else {
						log.Println("Block/Unblock action not set, skipping.")
					}

					if user_password != "" {
						User.SetPassword(user_password)
					} else {
						log.Println("Password not set, skipping.")
					}

					User.EmailChange()

					if user_role != "" {
						if !User.HasRole(user_role) {
							User.Roles = append(User.Roles, user_role)
							User.RolesAdd()
						}
					} else {
						log.Println("Role not set, skipping.")
					}

				} else {
					log.Printf("Could not find alias %v\n", Alias.GetName())
				}
			}
		}

		if !user_block && !user_create && !user_delete && !user_unblock && !user_verify {
			cmd.Usage()
			fmt.Println()
			log.Fatalln("no action specified.")
			os.Exit(1)
		}

	},
}

func init() {

	RootCmd.AddCommand(userCmd)

	userCmd.Flags().StringVarP(&user_name, "name", "n", "", "User name")
	userCmd.Flags().StringVarP(&user_email, "email", "e", "", "User email")
	userCmd.Flags().StringVarP(&user_password, "password", "s", "", "User password")
	userCmd.Flags().StringVarP(&user_role, "role", "r", "", "User role (case sensitive)")

	userCmd.Flags().StringVarP(&aliases, "aliases", "a", "", "Comma-separated list of drush aliases")
	userCmd.Flags().StringVarP(&pattern, "pattern", "p", "%v", "Pattern to match against drush aliases, where token is '%v'")

	userCmd.Flags().BoolVarP(&user_block, "block", "b", false, "Execute user block action.")
	userCmd.Flags().BoolVarP(&user_create, "create", "c", false, "Execute user create action.")
	userCmd.Flags().BoolVarP(&user_delete, "delete", "d", false, "Execute user delete action.")
	userCmd.Flags().BoolVarP(&user_unblock, "unblock", "u", false, "Execute user unblock action.")
	userCmd.Flags().BoolVarP(&user_verify, "verify", "v", false, "Execute user verification action.")

	userCmd.MarkFlagRequired("name")
	userCmd.MarkFlagRequired("email")
	userCmd.MarkFlagRequired("aliases")

}
<commit_msg>Prevent the re-execution of code requested behind user verification.<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"
	"net/mail"

	"os"
	"strings"

	"fmt"

	alias2 "github.com/fubarhouse/drubuild/alias"
	"github.com/fubarhouse/drubuild/user"
	"github.com/spf13/cobra"
)

// userCmd represents the user command
var userCmd = &cobra.Command{
	Use:   "user",
	Short: "User management tasks",
	Long:  ``,
	Run: func(cmd *cobra.Command, args []string) {

		if !strings.Contains(pattern, "%v") {
			log.Fatalf("Specified pattern '%v' does not include alias modifier '%%v'.", pattern)
			os.Exit(1)
		}

		Email, err := mail.ParseAddress(user_email)
		if err != nil {
			log.Fatal("Input email address could not be parsed.")
			log.Fatal(err)
		}

		if user_verify {
			for _, ThisAlias := range strings.Split(aliases, ",") {
				ThisAlias = strings.Trim(ThisAlias, " ")
				ThisAlias = strings.Replace(pattern, "%v", ThisAlias, 1)
				log.Printf("Beginning to work with alias %v.", ThisAlias)
				UserGroup := user.NewDrupalUserGroup()
				Alias := alias2.NewAlias(ThisAlias, "", ThisAlias)
				if Alias.GetStatus() {
					UserGroup.Populate(Alias.GetName())
					User := UserGroup.GetUser(user_name)
					if User.Alias == "" {
						User.Alias = Alias.GetName()
					}
					if User.Name == "" {
						User.Name = user_name
					}
					User.Email = Email.Address

					if user_create {
						User.Create(user_password)
					} else {
						log.Println("User creation not set, skipping.")
					}

					if user_block {
						User.State = 1
						User.StateChange()
					} else if user_unblock {
						User.State = 0
						User.StateChange()
					} else {
						log.Println("Block/Unblock action not set, skipping.")
					}

					if user_password != "" {
						User.SetPassword(user_password)
					} else {
						log.Println("Password not set, skipping.")
					}

					User.EmailChange()

					if user_role != "" {
						if !User.HasRole(user_role) {
							User.Roles = append(User.Roles, user_role)
							User.RolesAdd()
						}
					} else {
						log.Println("Role not set, skipping.")
					}

				} else {
					log.Printf("Could not find alias %v\n", Alias.GetName())
				}
			}
			return
		}

		if user_block {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				log.Printf("Beginning to work with alias %v.", Alias)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					if User.State == 1 {
						User.State = 0
						User.StateChange()
					} else {
						log.Printf("User '%v' is already blocked on %v\n", User.Name, Alias)
					}
				} else {
					log.Printf("User '%v' was not found on %v\n", User.Name, Alias)
				}
			}
			return
		}

		if user_create {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				log.Printf("Beginning to work with alias %v.", Alias)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				User.Alias = Alias
				User.Name = user_name
				User.Email = user_email
				User.Create(user_password)
				if user_role != "" {
					if !User.HasRole(user_role) {
						User.Roles = append(User.Roles, user_role)
					}
					User.RolesAdd()
				}
			}
			return
		}
		if user_delete {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				log.Printf("Beginning to work with alias %v.", Alias)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					User.Delete()
				} else {
					log.Printf("User '%v' was not found on %v", user_name, Alias)
				}
			}
			return
		}

		if user_unblock {
			for _, Alias := range strings.Split(aliases, ",") {
				Alias = strings.Trim(Alias, " ")
				Alias = strings.Replace(pattern, "%v", Alias, 1)
				log.Printf("Beginning to work with alias %v.", Alias)
				UserGroup := user.NewDrupalUserGroup()
				UserGroup.Populate(Alias)
				User := UserGroup.GetUser(user_name)
				if User.Name == user_name {
					if User.State == 0 {
						User.State = 1
						User.StateChange()
					} else {
						log.Printf("User '%v' is already unblocked on %v", User.Name, Alias)
					}
				} else {
					log.Printf("User '%v' was not found on %v", user_name, Alias)
				}
			}
			return
		}

		if !user_block && !user_create && !user_delete && !user_unblock && !user_verify {
			cmd.Usage()
			fmt.Println()
			log.Fatalln("no action specified.")
			os.Exit(1)
		}

	},
}

func init() {

	RootCmd.AddCommand(userCmd)

	userCmd.Flags().StringVarP(&user_name, "name", "n", "", "User name")
	userCmd.Flags().StringVarP(&user_email, "email", "e", "", "User email")
	userCmd.Flags().StringVarP(&user_password, "password", "s", "", "User password")
	userCmd.Flags().StringVarP(&user_role, "role", "r", "", "User role (case sensitive)")

	userCmd.Flags().StringVarP(&aliases, "aliases", "a", "", "Comma-separated list of drush aliases")
	userCmd.Flags().StringVarP(&pattern, "pattern", "p", "%v", "Pattern to match against drush aliases, where token is '%v'")

	userCmd.Flags().BoolVarP(&user_block, "block", "b", false, "Execute user block action.")
	userCmd.Flags().BoolVarP(&user_create, "create", "c", false, "Execute user create action.")
	userCmd.Flags().BoolVarP(&user_delete, "delete", "d", false, "Execute user delete action.")
	userCmd.Flags().BoolVarP(&user_unblock, "unblock", "u", false, "Execute user unblock action.")
	userCmd.Flags().BoolVarP(&user_verify, "verify", "v", false, "Execute user verification action.")

	userCmd.MarkFlagRequired("name")
	userCmd.MarkFlagRequired("email")
	userCmd.MarkFlagRequired("aliases")

}
<|endoftext|>


<commit_before>package gps

import (
	"fmt"
	"os"
	"os/exec"
	"testing"
	"time"
)

func mkTestCmd(iterations int) *monitoredCmd {
	return newMonitoredCmd(
		exec.Command("./echosleep", "-n", fmt.Sprint(iterations)),
		200*time.Millisecond,
	)
}

func TestMonitoredCmd(t *testing.T) {
	err := exec.Command("go", "build", "./_testdata/cmd/echosleep.go").Run()
	if err != nil {
		t.Errorf("Unable to build echosleep binary: %s", err)
	}
	defer os.Remove("./echosleep")

	cmd := mkTestCmd(2)
	err = cmd.run()
	if err != nil {
		t.Errorf("Expected command not to fail: %s", err)
	}

	expectedOutput := "foo\nfoo\n"
	if cmd.stdout.buf.String() != expectedOutput {
		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd.stdout.buf.String(), expectedOutput)
	}

	cmd2 := mkTestCmd(10)
	err = cmd2.run()
	if err == nil {
		t.Error("Expected command to fail")
	}

	_, ok := err.(*timeoutError)
	if !ok {
		t.Errorf("Expected a timeout error, but got: %s", err)
	}

	expectedOutput = "foo\nfoo\nfoo\nfoo\n"
	if cmd2.stdout.buf.String() != expectedOutput {
		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd2.stdout.buf.String(), expectedOutput)
	}
}
<commit_msg>Give tested monitored cmd a little slack<commit_after>package gps

import (
	"fmt"
	"os"
	"os/exec"
	"testing"
	"time"
)

func mkTestCmd(iterations int) *monitoredCmd {
	return newMonitoredCmd(
		exec.Command("./echosleep", "-n", fmt.Sprint(iterations)),
		(time.Duration(iterations+1))*100*time.Millisecond,
	)
}

func TestMonitoredCmd(t *testing.T) {
	err := exec.Command("go", "build", "./_testdata/cmd/echosleep.go").Run()
	if err != nil {
		t.Errorf("Unable to build echosleep binary: %s", err)
	}
	defer os.Remove("./echosleep")

	cmd := mkTestCmd(2)
	err = cmd.run()
	if err != nil {
		t.Errorf("Expected command not to fail: %s", err)
	}

	expectedOutput := "foo\nfoo\n"
	if cmd.stdout.buf.String() != expectedOutput {
		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd.stdout.buf.String(), expectedOutput)
	}

	cmd2 := mkTestCmd(10)
	err = cmd2.run()
	if err == nil {
		t.Error("Expected command to fail")
	}

	_, ok := err.(*timeoutError)
	if !ok {
		t.Errorf("Expected a timeout error, but got: %s", err)
	}

	expectedOutput = "foo\nfoo\nfoo\nfoo\n"
	if cmd2.stdout.buf.String() != expectedOutput {
		t.Errorf("Unexpected output:\n\t(GOT): %s\n\t(WNT): %s", cmd2.stdout.buf.String(), expectedOutput)
	}
}
<|endoftext|>


<commit_before>package driver

import (
	"github.com/realglobe-Inc/edo/util"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"path"
	"reflect"
	"time"
)

// mondodb を使うドライバー。

type mongoRegistry struct {
	dbName   string
	collName string
	*mgo.Session
}

func newMongoRegistry(url, dbName, collName string, indices []mgo.Index) (*mongoRegistry, error) {
	sess, err := mgo.Dial(url)
	if err != nil {
		return nil, erro.Wrap(err)
	}

	curIndices, err := sess.DB(dbName).C(collName).Indexes()
	if err != nil {
		return nil, erro.Wrap(err)
	}

	// 既存の要らない索引を消す。
	for _, curIdx := range curIndices {
		if len(curIdx.Key) == 1 && curIdx.Key[0] == "_id" {
			continue
		}

		ok := false
		for _, idx := range indices {
			if reflect.DeepEqual(curIdx, idx) {
				ok = true
				break
			}
		}
		if ok {
			continue
		}

		// 要らない。
		if err := sess.DB(dbName).C(collName).DropIndex(curIdx.Key...); err != nil {
			return nil, erro.Wrap(err)
		}
	}

	curIndices, err = sess.DB(dbName).C(collName).Indexes()
	if err != nil {
		return nil, erro.Wrap(err)
	}

	for _, idx := range indices {
		ok := true
		for _, curIdx := range curIndices {
			if reflect.DeepEqual(idx, curIdx) {
				ok = false
				break
			}
		}
		if !ok {
			// もうある。
			continue
		}

		if err := sess.DB(dbName).C(collName).EnsureIndex(idx); err != nil {
			return nil, erro.Wrap(err)
		}
	}

	return &mongoRegistry{dbName, collName, sess}, nil
}

// ログイン。
func NewMongoLoginRegistry(url, dbName, collName string) (LoginRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"access_token"},
			Unique:   true,
			DropDups: true,
			Sparse:   true,
		},
	})
}

type mongoUser struct {
	AccToken string `bson:"access_token"`

	UsrUuid string `bson:"user_uuid"`
}

func (reg *mongoRegistry) User(accToken string) (usrUuid string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"access_token": accToken})
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var res mongoUser
	if err := query.One(&res); err != nil {
		return "", erro.Wrap(err)
	}
	return res.UsrUuid, nil
}

// JavaScript.
func NewMongoJsRegistry(url, dbName, collName string) (JsRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"path"},
			Unique:   true,
			DropDups: true,
			Sparse:   true,
		},
	})
}

type mongoObject struct {
	Path string `bson:"path"`

	Service bool     `bson:"service,omitempty"`
	Library bool     `bson:"library,omitempty"`
	Include []string `bson:"include,omitempty"`
	Code    string   `bson:"code"`

	Date   time.Time `bson:"date"`
	Digest int       `bson:"digest"`
}

func (reg *mongoRegistry) Object(dir, objName string) (*Object, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"path": path.Join(dir, objName)})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoObject
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return &Object{res.Service, res.Library, res.Include, res.Code}, nil
}

func (obj *Object) digest() int {
	prime := 31
	dig := 0
	dig = prime*dig + util.DigestBool(obj.Service)
	dig = prime*dig + util.DigestBool(obj.Library)
	for _, inc := range obj.Include {
		dig = prime*dig + util.DigestString(inc)
	}
	dig = prime*dig + util.DigestString(obj.Code)
	return dig

}

func (reg *mongoRegistry) AddObject(dir, objName string, obj *Object) error {
	mongoObj := &mongoObject{path.Join(dir, objName), obj.Service, obj.Library, obj.Include, obj.Code, time.Now(), obj.digest()}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"path": mongoObj.Path}, mongoObj); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveObject(dir, objName string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"path": path.Join(dir, objName)}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ユーザー情報。
func NewMongoUserRegistry(url, dbName, collName string) (UserRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:    []string{"user_uuid"},
			Sparse: true,
		},
		mgo.Index{
			Key:      []string{"user_uuid", "key"},
			Unique:   true,
			DropDups: true,
			Sparse:   true,
		},
	})
}

type mongoAttribute struct {
	UsrUuid  string      `bson:"user_uuid"`
	AttrName string      `bson:"key"`
	Attr     interface{} `bson:"value"`
}

func (reg *mongoRegistry) Attributes(usrUuid string) (attrs map[string]interface{}, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	mongoAttrs := []mongoAttribute{}
	if err := query.Iter().All(&mongoAttrs); err != nil {
		return nil, erro.Wrap(err)
	}
	attrs = map[string]interface{}{}
	for _, mongoAttr := range mongoAttrs {
		attrs[mongoAttr.AttrName] = mongoAttr.Attr
	}
	return attrs, nil
}

func (reg *mongoRegistry) Attribute(usrUuid, attrName string) (attr interface{}, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid, "key": attrName})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var mongoAttr mongoAttribute
	if err := query.One(&mongoAttr); err != nil {
		return nil, erro.Wrap(err)
	}
	return mongoAttr.Attr, nil
}

func (reg *mongoRegistry) AddAttribute(usrUuid, attrName string, attr interface{}) error {
	mongoAttr := &mongoAttribute{usrUuid, attrName, attr}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"user_uuid": usrUuid, "key": attrName}, mongoAttr); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveAttribute(usrUuid, attrName string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"user_uuid": usrUuid, "key": attrName}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ジョブ。
func NewMongoJobRegistry(url, dbName, collName string) (JobRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"job_id"},
			Unique:   true,
			DropDups: true,
			Sparse:   true,
		},
		mgo.Index{
			Key: []string{"deadline"},
		},
	})
}

type mongoJobResult struct {
	JobId    string            `bson:"job_id"`
	Deadline time.Time         `bson:"deadline"`
	Status   int               `bson:"status"`
	Headers  map[string]string `bson:"headers,omitempty"`
	Body     string            `bson:"body,omitempty"`
}

func (reg *mongoRegistry) Result(jobId string) (*JobResult, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"job_id": jobId})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoJobResult
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return &JobResult{res.Status, res.Headers, res.Body}, nil
}

func (reg *mongoRegistry) AddResult(jobId string, res *JobResult, deadline time.Time) error {
	mongoRes := &mongoJobResult{jobId, deadline, res.Status, res.Headers, res.Body}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"job_id": jobId}, mongoRes); err != nil {
		return erro.Wrap(err)
	}

	// 古いのを削除。
	reg.DB(reg.dbName).C(reg.collName).RemoveAll(bson.M{"deadline": bson.M{"$lt": time.Now()}})

	return nil
}

// 別名。
func NewMongoNameRegistry(url, dbName, collName string) (NameRegistry, error) {
	sess, err := mgo.Dial(url)
	if err != nil {
		return nil, erro.Wrap(err)
	}

	nameIdx := mgo.Index{
		Key:      []string{"name"},
		Unique:   true,
		DropDups: true,
		Sparse:   true,
	}
	if err := sess.DB(dbName).C(collName).EnsureIndex(nameIdx); err != nil {
		return nil, erro.Wrap(err)
	}

	return &mongoRegistry{dbName, collName, sess}, nil
}

type mongoAddress struct {
	Name string `bson:"name"`
	Addr string `bson:"address"`
}

func (reg *mongoRegistry) Address(name string) (addr string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"name": name})
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var mongoAddr mongoAddress
	if err := query.One(&mongoAddr); err != nil {
		return "", erro.Wrap(err)
	}
	return mongoAddr.Addr, nil
}

func (reg *mongoRegistry) Addresses(name string) (addrs []string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"name": bson.M{"$regex": name + "$"}})
	var mongoAddrs []mongoAddress
	if err := query.All(&mongoAddrs); err != nil {
		return nil, erro.Wrap(err)
	}
	for _, mongoAddr := range mongoAddrs {
		addrs = append(addrs, mongoAddr.Addr)
	}
	return addrs, nil
}

// イベント。
func NewMongoEventRegistry(url, dbName, collName string) (EventRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"user_uuid", "event"},
			Unique:   true,
			DropDups: true,
			Sparse:   true,
		},
	})
}

type mongoHandler struct {
	UsrUuid string  `bson:"user_uuid"`
	Event   string  `bson:"event"`
	Hndl    Handler `bson:"handler"`
}

func (reg *mongoRegistry) Handler(usrUuid, event string) (Handler, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid, "event": event})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoHandler
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return res.Hndl, nil
}

func (reg *mongoRegistry) AddHandler(usrUuid, event string, hndl Handler) error {
	mongoHndl := &mongoHandler{usrUuid, event, hndl}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"user_uuid": usrUuid, "event": event}, mongoHndl); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveHandler(usrUuid, event string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"user_uuid": usrUuid, "event": event}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ログイン。
// エンドポイントごとに保存しても親を辿るのが面倒なので 1 エントリで。
// TODO エンドポイントごとにして、任意の親を列挙する方法があるか？
func NewMongoServiceRegistry(url, dbName, collName string) (ServiceRegistry, error) {
	return newMongoRegistry(url, dbName, collName, nil)
}

type mongoService struct {
	Cont map[string]string `bson:"container"`
}

func (reg *mongoRegistry) Service(endPt string) (usrUuid string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(nil)
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var res mongoService
	if err := query.One(&res); err != nil {
		return "", erro.Wrap(err)
	}

	tree := newServiceTree()
	tree.fromContainer(res.Cont)

	return tree.service(endPt), nil
}
<commit_msg>全然 Sparse ではなかったので、mongodb のインデクス作成時に Sparse を指定しないようにした<commit_after>package driver

import (
	"github.com/realglobe-Inc/edo/util"
	"github.com/realglobe-Inc/go-lib-rg/erro"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"path"
	"reflect"
	"time"
)

// mondodb を使うドライバー。

type mongoRegistry struct {
	dbName   string
	collName string
	*mgo.Session
}

func newMongoRegistry(url, dbName, collName string, indices []mgo.Index) (*mongoRegistry, error) {
	sess, err := mgo.Dial(url)
	if err != nil {
		return nil, erro.Wrap(err)
	}

	curIndices, err := sess.DB(dbName).C(collName).Indexes()
	if err != nil {
		return nil, erro.Wrap(err)
	}

	// 既存の要らない索引を消す。
	for _, curIdx := range curIndices {
		if len(curIdx.Key) == 1 && curIdx.Key[0] == "_id" {
			continue
		}

		ok := false
		for _, idx := range indices {
			if reflect.DeepEqual(curIdx, idx) {
				ok = true
				break
			}
		}
		if ok {
			continue
		}

		// 要らない。
		if err := sess.DB(dbName).C(collName).DropIndex(curIdx.Key...); err != nil {
			return nil, erro.Wrap(err)
		}
	}

	curIndices, err = sess.DB(dbName).C(collName).Indexes()
	if err != nil {
		return nil, erro.Wrap(err)
	}

	for _, idx := range indices {
		ok := true
		for _, curIdx := range curIndices {
			if reflect.DeepEqual(idx, curIdx) {
				ok = false
				break
			}
		}
		if !ok {
			// もうある。
			continue
		}

		if err := sess.DB(dbName).C(collName).EnsureIndex(idx); err != nil {
			return nil, erro.Wrap(err)
		}
	}

	return &mongoRegistry{dbName, collName, sess}, nil
}

// ログイン。
func NewMongoLoginRegistry(url, dbName, collName string) (LoginRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"access_token"},
			Unique:   true,
			DropDups: true,
		},
	})
}

type mongoUser struct {
	AccToken string `bson:"access_token"`

	UsrUuid string `bson:"user_uuid"`
}

func (reg *mongoRegistry) User(accToken string) (usrUuid string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"access_token": accToken})
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var res mongoUser
	if err := query.One(&res); err != nil {
		return "", erro.Wrap(err)
	}
	return res.UsrUuid, nil
}

// JavaScript.
func NewMongoJsRegistry(url, dbName, collName string) (JsRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"path"},
			Unique:   true,
			DropDups: true,
		},
	})
}

type mongoObject struct {
	Path string `bson:"path"`

	Service bool     `bson:"service,omitempty"`
	Library bool     `bson:"library,omitempty"`
	Include []string `bson:"include,omitempty"`
	Code    string   `bson:"code"`

	Date   time.Time `bson:"date"`
	Digest int       `bson:"digest"`
}

func (reg *mongoRegistry) Object(dir, objName string) (*Object, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"path": path.Join(dir, objName)})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoObject
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return &Object{res.Service, res.Library, res.Include, res.Code}, nil
}

func (obj *Object) digest() int {
	prime := 31
	dig := 0
	dig = prime*dig + util.DigestBool(obj.Service)
	dig = prime*dig + util.DigestBool(obj.Library)
	for _, inc := range obj.Include {
		dig = prime*dig + util.DigestString(inc)
	}
	dig = prime*dig + util.DigestString(obj.Code)
	return dig

}

func (reg *mongoRegistry) AddObject(dir, objName string, obj *Object) error {
	mongoObj := &mongoObject{path.Join(dir, objName), obj.Service, obj.Library, obj.Include, obj.Code, time.Now(), obj.digest()}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"path": mongoObj.Path}, mongoObj); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveObject(dir, objName string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"path": path.Join(dir, objName)}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ユーザー情報。
func NewMongoUserRegistry(url, dbName, collName string) (UserRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key: []string{"user_uuid"},
		},
		mgo.Index{
			Key:      []string{"user_uuid", "key"},
			Unique:   true,
			DropDups: true,
		},
	})
}

type mongoAttribute struct {
	UsrUuid  string      `bson:"user_uuid"`
	AttrName string      `bson:"key"`
	Attr     interface{} `bson:"value"`
}

func (reg *mongoRegistry) Attributes(usrUuid string) (attrs map[string]interface{}, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	mongoAttrs := []mongoAttribute{}
	if err := query.Iter().All(&mongoAttrs); err != nil {
		return nil, erro.Wrap(err)
	}
	attrs = map[string]interface{}{}
	for _, mongoAttr := range mongoAttrs {
		attrs[mongoAttr.AttrName] = mongoAttr.Attr
	}
	return attrs, nil
}

func (reg *mongoRegistry) Attribute(usrUuid, attrName string) (attr interface{}, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid, "key": attrName})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var mongoAttr mongoAttribute
	if err := query.One(&mongoAttr); err != nil {
		return nil, erro.Wrap(err)
	}
	return mongoAttr.Attr, nil
}

func (reg *mongoRegistry) AddAttribute(usrUuid, attrName string, attr interface{}) error {
	mongoAttr := &mongoAttribute{usrUuid, attrName, attr}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"user_uuid": usrUuid, "key": attrName}, mongoAttr); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveAttribute(usrUuid, attrName string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"user_uuid": usrUuid, "key": attrName}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ジョブ。
func NewMongoJobRegistry(url, dbName, collName string) (JobRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"job_id"},
			Unique:   true,
			DropDups: true,
		},
		mgo.Index{
			Key: []string{"deadline"},
		},
	})
}

type mongoJobResult struct {
	JobId    string            `bson:"job_id"`
	Deadline time.Time         `bson:"deadline"`
	Status   int               `bson:"status"`
	Headers  map[string]string `bson:"headers,omitempty"`
	Body     string            `bson:"body,omitempty"`
}

func (reg *mongoRegistry) Result(jobId string) (*JobResult, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"job_id": jobId})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoJobResult
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return &JobResult{res.Status, res.Headers, res.Body}, nil
}

func (reg *mongoRegistry) AddResult(jobId string, res *JobResult, deadline time.Time) error {
	mongoRes := &mongoJobResult{jobId, deadline, res.Status, res.Headers, res.Body}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"job_id": jobId}, mongoRes); err != nil {
		return erro.Wrap(err)
	}

	// 古いのを削除。
	reg.DB(reg.dbName).C(reg.collName).RemoveAll(bson.M{"deadline": bson.M{"$lt": time.Now()}})

	return nil
}

// 別名。
func NewMongoNameRegistry(url, dbName, collName string) (NameRegistry, error) {
	sess, err := mgo.Dial(url)
	if err != nil {
		return nil, erro.Wrap(err)
	}

	nameIdx := mgo.Index{
		Key:      []string{"name"},
		Unique:   true,
		DropDups: true,
	}
	if err := sess.DB(dbName).C(collName).EnsureIndex(nameIdx); err != nil {
		return nil, erro.Wrap(err)
	}

	return &mongoRegistry{dbName, collName, sess}, nil
}

type mongoAddress struct {
	Name string `bson:"name"`
	Addr string `bson:"address"`
}

func (reg *mongoRegistry) Address(name string) (addr string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"name": name})
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var mongoAddr mongoAddress
	if err := query.One(&mongoAddr); err != nil {
		return "", erro.Wrap(err)
	}
	return mongoAddr.Addr, nil
}

func (reg *mongoRegistry) Addresses(name string) (addrs []string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"name": bson.M{"$regex": name + "$"}})
	var mongoAddrs []mongoAddress
	if err := query.All(&mongoAddrs); err != nil {
		return nil, erro.Wrap(err)
	}
	for _, mongoAddr := range mongoAddrs {
		addrs = append(addrs, mongoAddr.Addr)
	}
	return addrs, nil
}

// イベント。
func NewMongoEventRegistry(url, dbName, collName string) (EventRegistry, error) {
	return newMongoRegistry(url, dbName, collName, []mgo.Index{
		mgo.Index{
			Key:      []string{"user_uuid", "event"},
			Unique:   true,
			DropDups: true,
		},
	})
}

type mongoHandler struct {
	UsrUuid string  `bson:"user_uuid"`
	Event   string  `bson:"event"`
	Hndl    Handler `bson:"handler"`
}

func (reg *mongoRegistry) Handler(usrUuid, event string) (Handler, error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(bson.M{"user_uuid": usrUuid, "event": event})
	if n, err := query.Count(); err != nil {
		return nil, erro.Wrap(err)
	} else if n == 0 {
		return nil, nil
	}
	var res mongoHandler
	if err := query.One(&res); err != nil {
		return nil, erro.Wrap(err)
	}
	return res.Hndl, nil
}

func (reg *mongoRegistry) AddHandler(usrUuid, event string, hndl Handler) error {
	mongoHndl := &mongoHandler{usrUuid, event, hndl}
	if _, err := reg.DB(reg.dbName).C(reg.collName).Upsert(bson.M{"user_uuid": usrUuid, "event": event}, mongoHndl); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

func (reg *mongoRegistry) RemoveHandler(usrUuid, event string) error {
	if err := reg.DB(reg.dbName).C(reg.collName).Remove(bson.M{"user_uuid": usrUuid, "event": event}); err != nil {
		return erro.Wrap(err)
	}
	return nil
}

// ログイン。
// エンドポイントごとに保存しても親を辿るのが面倒なので 1 エントリで。
// TODO エンドポイントごとにして、任意の親を列挙する方法があるか？
func NewMongoServiceRegistry(url, dbName, collName string) (ServiceRegistry, error) {
	return newMongoRegistry(url, dbName, collName, nil)
}

type mongoService struct {
	Cont map[string]string `bson:"container"`
}

func (reg *mongoRegistry) Service(endPt string) (usrUuid string, err error) {
	query := reg.DB(reg.dbName).C(reg.collName).Find(nil)
	if n, err := query.Count(); err != nil {
		return "", erro.Wrap(err)
	} else if n == 0 {
		return "", nil
	}
	var res mongoService
	if err := query.One(&res); err != nil {
		return "", erro.Wrap(err)
	}

	tree := newServiceTree()
	tree.fromContainer(res.Cont)

	return tree.service(endPt), nil
}
<|endoftext|>


<commit_before>package build

func (b *Build) Linux() error {
	return nil
}
<commit_msg>build: first draft of deb package building<commit_after>package build

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"go/build"
	"io"
	"os"
	"strings"
	"time"

	"github.com/blakesmith/ar"
)

const control = `Package: %s
Version: %s
Architecture: %s
Maintainer: Koding Developers <hello@koding.com>
Installed-Size: %d
Section: devel
Priority: extra
Description: %s Kite
`

func (b *Build) Linux() error {
	debName := b.Output + ".deb"
	tarFile := b.Output + ".tar.gz"

	deb, err := os.Create(debName + ".inprogress")
	if err != nil {
		return fmt.Errorf("cannot create deb: %v", err)
	}
	defer deb.Close()

	// create first tar file
	fmt.Println("creating tarfile", tarFile)
	if err := b.TarGzFile(); err != nil {
		return err
	}

	fmt.Println("openening tarfile", tarFile)
	tf, err := os.Open(tarFile)
	if err != nil {
		return err
	}
	defer tf.Close()

	fmt.Println("creating debfile from tarfile")
	if err := b.createDeb(tf, deb); err != nil {
		return err
	}

	if err := os.Rename(debName+".inprogress", debName); err != nil {
		return err
	}

	fmt.Println("package", debName, "ready")
	return nil
}

func debArch() string {
	arch := build.Default.GOARCH
	if arch == "386" {
		return "i386"
	}
	return arch
}

func (b *Build) createDeb(tarball io.Reader, deb io.Writer) error {
	now := time.Now()
	dataTarGz, md5sums, instSize, err := b.translateTarball(now, tarball)
	if err != nil {
		return err
	}

	controlTarGz, err := b.createControl(now, instSize, md5sums)
	if err != nil {
		return err
	}
	w := ar.NewWriter(deb)
	if err := w.WriteGlobalHeader(); err != nil {
		return fmt.Errorf("cannot write ar header to deb file: %v", err)
	}

	if err := addArFile(now, w, "debian-binary", []byte("2.0\n")); err != nil {
		return fmt.Errorf("cannot pack debian-binary: %v", err)
	}
	if err := addArFile(now, w, "control.tar.gz", controlTarGz); err != nil {
		return fmt.Errorf("cannot add control.tar.gz to deb: %v", err)
	}
	if err := addArFile(now, w, "data.tar.gz", dataTarGz); err != nil {
		return fmt.Errorf("cannot add data.tar.gz to deb: %v", err)
	}
	return nil
}

func addArFile(now time.Time, w *ar.Writer, name string, body []byte) error {
	hdr := ar.Header{
		Name:    name,
		Size:    int64(len(body)),
		Mode:    0644,
		ModTime: now,
	}
	if err := w.WriteHeader(&hdr); err != nil {
		return fmt.Errorf("cannot write file header: %v", err)
	}
	_, err := w.Write(body)
	return err
}

func (b *Build) translateTarball(now time.Time, tarball io.Reader) (dataTarGz, md5sums []byte, instSize int64, err error) {
	buf := &bytes.Buffer{}
	compress := gzip.NewWriter(buf)
	out := tar.NewWriter(compress)

	md5buf := &bytes.Buffer{}
	md5tmp := make([]byte, 0, md5.Size)

	uncompress, err := gzip.NewReader(tarball)
	if err != nil {
		return nil, nil, 0, fmt.Errorf("cannot uncompress tarball: %v", err)
	}

	in := tar.NewReader(uncompress)
	first := true
	for {
		h, err := in.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return nil, nil, 0, fmt.Errorf("cannot read tarball: %v", err)
		}

		instSize += h.Size
		h.Name = strings.TrimLeft(h.Name, "./")
		if first && h.Name != b.AppName && h.Name != b.AppName+"/" {
			h := tar.Header{
				Name:     "./opt/koding/" + b.AppName,
				Mode:     0755,
				ModTime:  h.ModTime,
				Typeflag: tar.TypeDir,
			}

			if err := out.WriteHeader(&h); err != nil {
				return nil, nil, 0, fmt.Errorf("cannot write header of %s to data.tar.gz: %v", h.Name, err)
			}
		}

		const prefix = "./opt/koding/"
		h.Name = prefix + h.Name
		if h.Typeflag == tar.TypeDir && !strings.HasSuffix(h.Name, "/") {
			h.Name += "/"
		}

		if err := out.WriteHeader(h); err != nil {
			return nil, nil, 0, fmt.Errorf("cannot write header of %s to data.tar.gz: %v", h.Name, err)
		}

		fmt.Println("tar: packing", h.Name[len(prefix):])

		if h.Typeflag == tar.TypeDir {
			continue
		}

		digest := md5.New()
		if _, err := io.Copy(out, io.TeeReader(in, digest)); err != nil {
			return nil, nil, 0, err
		}

		fmt.Fprintf(md5buf, "%x  %s\n", digest.Sum(md5tmp), h.Name[2:])
	}

	if err := out.Close(); err != nil {
		return nil, nil, 0, err
	}

	if err := compress.Close(); err != nil {
		return nil, nil, 0, err
	}

	return buf.Bytes(), md5buf.Bytes(), instSize, nil
}

func addTarSymlink(now time.Time, out *tar.Writer, name, target string) error {
	h := tar.Header{
		Name:     name,
		Linkname: target,
		Mode:     0777,
		ModTime:  now,
		Typeflag: tar.TypeSymlink,
	}
	if err := out.WriteHeader(&h); err != nil {
		return fmt.Errorf("cannot write header of %s to data.tar.gz: %v", h.Name, err)
	}
	return nil
}

func (b *Build) createControl(now time.Time, instSize int64, md5sums []byte) (controlTarGz []byte, err error) {
	buf := &bytes.Buffer{}
	compress := gzip.NewWriter(buf)
	tarball := tar.NewWriter(compress)

	body := []byte(fmt.Sprintf(control, b.Version, debArch(), instSize/1024))
	hdr := tar.Header{
		Name:     "control",
		Size:     int64(len(body)),
		Mode:     0644,
		ModTime:  now,
		Typeflag: tar.TypeReg,
	}
	if err := tarball.WriteHeader(&hdr); err != nil {
		return nil, fmt.Errorf("cannot write header of control file to control.tar.gz: %v", err)
	}
	if _, err := tarball.Write(body); err != nil {
		return nil, fmt.Errorf("cannot write control file to control.tar.gz: %v", err)
	}

	hdr = tar.Header{
		Name:     "md5sums",
		Size:     int64(len(md5sums)),
		Mode:     0644,
		ModTime:  now,
		Typeflag: tar.TypeReg,
	}
	if err := tarball.WriteHeader(&hdr); err != nil {
		return nil, fmt.Errorf("cannot write header of md5sums file to control.tar.gz: %v", err)
	}
	if _, err := tarball.Write(md5sums); err != nil {
		return nil, fmt.Errorf("cannot write md5sums file to control.tar.gz: %v", err)
	}

	if err := tarball.Close(); err != nil {
		return nil, fmt.Errorf("closing control.tar.gz: %v", err)
	}
	if err := compress.Close(); err != nil {
		return nil, fmt.Errorf("closing control.tar.gz: %v", err)
	}
	return buf.Bytes(), nil
}
<|endoftext|>


<commit_before>package commands

import (
	"context"
	"fmt"
	"io/ioutil"

	"github.com/killbill/kbcli/kbclient/catalog"
	"github.com/killbill/kbcli/kbcmd/cmdlib"
	"github.com/urfave/cli"
)

func uploadCatalog(ctx context.Context, o *cmdlib.Options) error {
	if len(o.Args) != 1 {
		return cmdlib.ErrorInvalidArgs
	}
	catalogFile := o.Args[0]
	contents, err := ioutil.ReadFile(catalogFile)
	if err != nil {
		return fmt.Errorf("unable to read catalog file %s. %v", catalogFile, err)
	}
	strContents := string(contents)

	_, err = o.Client().Catalog.UploadCatalogXML(ctx, &catalog.UploadCatalogXMLParams{
		Body: &strContents,
	})
	if err != nil {
		return err
	}
	o.Outputln("catalog successfully uploaded")
	return err
}

func registerCatalogCommands(r *cmdlib.App) {
	// Register top level command
	r.Register("", cli.Command{
		Name:    "catalog",
		Aliases: []string{"cat"},
		Usage:   "catalog related commands",
	}, nil)

	r.Register("catalog", cli.Command{
		Name:        "upload",
		Usage:       "upload given catalog",
		ArgsUsage:   "CATALOG_FILE",
		Description: "Uploads the catalog",
	}, uploadCatalog)
}
<commit_msg>Add catalog get command<commit_after>package commands

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"

	"github.com/killbill/kbcli/kbclient/catalog"
	"github.com/killbill/kbcli/kbcmd/cmdlib"
	"github.com/urfave/cli"
)

func getCatalog(ctx context.Context, o *cmdlib.Options) error {
	if len(o.Args) > 1 {
		return cmdlib.ErrorInvalidArgs
	}
	var outputFile string

	if len(o.Args) == 1 {
		outputFile = o.Args[0]
		_, err := os.Stat(outputFile)
		if err == nil || !os.IsNotExist(err) {
			return fmt.Errorf("output file already exists %v", outputFile)
		}
	}

	res, err := o.Client().Catalog.GetCatalogXML(ctx, &catalog.GetCatalogXMLParams{})
	if err != nil {
		return err
	}

	if outputFile != "" {
		if err = ioutil.WriteFile(outputFile, []byte(res.Payload), 0644); err != nil {
			return fmt.Errorf("unable to write output: %v", err)
		}
		o.Outputln("catalog successfully written to %v", outputFile)
	} else {
		o.Outputln("%s", res.Payload)
	}
	return nil
}

func uploadCatalog(ctx context.Context, o *cmdlib.Options) error {
	if len(o.Args) != 1 {
		return cmdlib.ErrorInvalidArgs
	}
	catalogFile := o.Args[0]
	contents, err := ioutil.ReadFile(catalogFile)
	if err != nil {
		return fmt.Errorf("unable to read catalog file %s. %v", catalogFile, err)
	}
	strContents := string(contents)

	_, err = o.Client().Catalog.UploadCatalogXML(ctx, &catalog.UploadCatalogXMLParams{
		Body: &strContents,
	})
	if err != nil {
		return err
	}
	o.Outputln("catalog successfully uploaded")
	return err
}

func registerCatalogCommands(r *cmdlib.App) {
	// Register top level command
	r.Register("", cli.Command{
		Name:    "catalog",
		Aliases: []string{"cat"},
		Usage:   "catalog related commands",
	}, nil)

	r.Register("catalog", cli.Command{
		Name:  "get",
		Usage: "get catalog for tenant",
		ArgsUsage: `[OUTPUT_FILE]

		For ex.,

		# To write the catalog to a file
		kbcmd catalog upload /tmp/catalog.json

		# To write to stdout
		kbcmd catalog upload
`,
		Description: "Retrieves the catalog and stores in the file",
	}, getCatalog)

	r.Register("catalog", cli.Command{
		Name:  "upload",
		Usage: "upload given catalog",
		ArgsUsage: `CATALOG_FILE

For ex.,
		kbcmd catalog upload /tmp/catalog.json
`,
		Description: "Uploads the catalog",
	}, uploadCatalog)
}
<|endoftext|>


<commit_before>package repository

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/restic/restic/backend"
	"github.com/restic/restic/debug"
	"github.com/restic/restic/pack"
)

// Index holds a lookup table for id -> pack.
type Index struct {
	m    sync.Mutex
	pack map[string]indexEntry
}

type indexEntry struct {
	tpe    pack.BlobType
	packID backend.ID
	offset uint
	length uint
	old    bool
}

// NewIndex returns a new index.
func NewIndex() *Index {
	return &Index{
		pack: make(map[string]indexEntry),
	}
}

func (idx *Index) store(t pack.BlobType, id, pack backend.ID, offset, length uint, old bool) {
	idx.pack[id.String()] = indexEntry{
		tpe:    t,
		packID: pack,
		offset: offset,
		length: length,
		old:    old,
	}
}

// Store remembers the id and pack in the index.
func (idx *Index) Store(t pack.BlobType, id, pack backend.ID, offset, length uint) {
	idx.m.Lock()
	defer idx.m.Unlock()

	debug.Log("Index.Store", "pack %v contains id %v (%v), offset %v, length %v",
		pack.Str(), id.Str(), t, offset, length)

	idx.store(t, id, pack, offset, length, false)
}

// Remove removes the pack ID from the index.
func (idx *Index) Remove(packID backend.ID) {
	idx.m.Lock()
	defer idx.m.Unlock()

	debug.Log("Index.Remove", "id %v removed", packID.Str())

	s := packID.String()
	if _, ok := idx.pack[s]; ok {
		delete(idx.pack, s)
	}
}

// Lookup returns the pack for the id.
func (idx *Index) Lookup(id backend.ID) (packID backend.ID, tpe pack.BlobType, offset, length uint, err error) {
	idx.m.Lock()
	defer idx.m.Unlock()

	if p, ok := idx.pack[id.String()]; ok {
		debug.Log("Index.Lookup", "id %v found in pack %v at %d, length %d",
			id.Str(), p.packID.Str(), p.offset, p.length)
		return p.packID, p.tpe, p.offset, p.length, nil
	}

	debug.Log("Index.Lookup", "id %v not found", id.Str())
	return nil, pack.Data, 0, 0, errors.New("id not found")
}

// Has returns true iff the id is listed in the index.
func (idx *Index) Has(id backend.ID) bool {
	_, _, _, _, err := idx.Lookup(id)
	if err == nil {
		return true
	}

	return false
}

// Merge loads all items from other into idx.
func (idx *Index) Merge(other *Index) {
	debug.Log("Index.Merge", "Merge index with %p", other)
	idx.m.Lock()
	defer idx.m.Unlock()

	for k, v := range other.pack {
		if _, ok := idx.pack[k]; ok {
			debug.Log("Index.Merge", "index already has key %v, updating", k[:8])
		}

		idx.pack[k] = v
	}
	debug.Log("Index.Merge", "done merging index")
}

// Each returns a channel that yields all blobs known to the index. If done is
// closed, the background goroutine terminates. This blocks any modification of
// the index.
func (idx *Index) Each(done chan struct{}) <-chan pack.Blob {
	idx.m.Lock()

	ch := make(chan pack.Blob)

	go func() {
		defer idx.m.Unlock()
		defer func() {
			close(ch)
		}()

		for ids, blob := range idx.pack {
			id, err := backend.ParseID(ids)
			if err != nil {
				// ignore invalid IDs
				continue
			}

			select {
			case <-done:
				return
			case ch <- pack.Blob{
				ID:     id,
				Offset: blob.offset,
				Type:   blob.tpe,
				Length: uint32(blob.length),
			}:
			}
		}
	}()

	return ch
}

// Count returns the number of blobs of type t in the index.
func (idx *Index) Count(t pack.BlobType) (n uint) {
	debug.Log("Index.Count", "counting blobs of type %v", t)
	idx.m.Lock()
	defer idx.m.Unlock()

	for id, blob := range idx.pack {
		if blob.tpe == t {
			n++
			debug.Log("Index.Count", "  blob %v counted: %v", id[:8], blob)
		}
	}

	return
}

type packJSON struct {
	ID    string     `json:"id"`
	Blobs []blobJSON `json:"blobs"`
}

type blobJSON struct {
	ID     string        `json:"id"`
	Type   pack.BlobType `json:"type"`
	Offset uint          `json:"offset"`
	Length uint          `json:"length"`
}

// Encode writes the JSON serialization of the index to the writer w. This
// serialization only contains new blobs added via idx.Store(), not old ones
// introduced via DecodeIndex().
func (idx *Index) Encode(w io.Writer) error {
	debug.Log("Index.Encode", "encoding index")
	idx.m.Lock()
	defer idx.m.Unlock()

	list := []*packJSON{}
	packs := make(map[string]*packJSON)

	for id, blob := range idx.pack {
		if blob.old {
			continue
		}

		debug.Log("Index.Encode", "handle blob %q", id[:8])

		if blob.packID == nil {
			debug.Log("Index.Encode", "blob %q has no packID! (type %v, offset %v, length %v)",
				id[:8], blob.tpe, blob.offset, blob.length)
			return fmt.Errorf("unable to serialize index: pack for blob %v hasn't been written yet", id)
		}

		// see if pack is already in map
		p, ok := packs[blob.packID.String()]
		if !ok {
			// else create new pack
			p = &packJSON{ID: blob.packID.String()}

			// and append it to the list and map
			list = append(list, p)
			packs[p.ID] = p
		}

		// add blob
		p.Blobs = append(p.Blobs, blobJSON{
			ID:     id,
			Type:   blob.tpe,
			Offset: blob.offset,
			Length: blob.length,
		})
	}

	debug.Log("Index.Encode", "done")

	enc := json.NewEncoder(w)
	return enc.Encode(list)
}

// DecodeIndex loads and unserializes an index from rd.
func DecodeIndex(rd io.Reader) (*Index, error) {
	debug.Log("Index.DecodeIndex", "Start decoding index")
	list := []*packJSON{}

	dec := json.NewDecoder(rd)
	err := dec.Decode(&list)
	if err != nil {
		return nil, err
	}

	idx := NewIndex()
	for _, pack := range list {
		packID, err := backend.ParseID(pack.ID)
		if err != nil {
			debug.Log("Index.DecodeIndex", "error parsing pack ID %q: %v", pack.ID, err)
			return nil, err
		}

		for _, blob := range pack.Blobs {
			blobID, err := backend.ParseID(blob.ID)
			if err != nil {
				debug.Log("Index.DecodeIndex", "error parsing blob ID %q: %v", blob.ID, err)
				return nil, err
			}

			idx.store(blob.Type, blobID, packID, blob.Offset, blob.Length, true)
		}
	}

	debug.Log("Index.DecodeIndex", "done")
	return idx, err
}
<commit_msg>Add index.Dump()<commit_after>package repository

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/restic/restic/backend"
	"github.com/restic/restic/debug"
	"github.com/restic/restic/pack"
)

// Index holds a lookup table for id -> pack.
type Index struct {
	m    sync.Mutex
	pack map[string]indexEntry
}

type indexEntry struct {
	tpe    pack.BlobType
	packID backend.ID
	offset uint
	length uint
	old    bool
}

// NewIndex returns a new index.
func NewIndex() *Index {
	return &Index{
		pack: make(map[string]indexEntry),
	}
}

func (idx *Index) store(t pack.BlobType, id, pack backend.ID, offset, length uint, old bool) {
	idx.pack[id.String()] = indexEntry{
		tpe:    t,
		packID: pack,
		offset: offset,
		length: length,
		old:    old,
	}
}

// Store remembers the id and pack in the index.
func (idx *Index) Store(t pack.BlobType, id, pack backend.ID, offset, length uint) {
	idx.m.Lock()
	defer idx.m.Unlock()

	debug.Log("Index.Store", "pack %v contains id %v (%v), offset %v, length %v",
		pack.Str(), id.Str(), t, offset, length)

	idx.store(t, id, pack, offset, length, false)
}

// Remove removes the pack ID from the index.
func (idx *Index) Remove(packID backend.ID) {
	idx.m.Lock()
	defer idx.m.Unlock()

	debug.Log("Index.Remove", "id %v removed", packID.Str())

	s := packID.String()
	if _, ok := idx.pack[s]; ok {
		delete(idx.pack, s)
	}
}

// Lookup returns the pack for the id.
func (idx *Index) Lookup(id backend.ID) (packID backend.ID, tpe pack.BlobType, offset, length uint, err error) {
	idx.m.Lock()
	defer idx.m.Unlock()

	if p, ok := idx.pack[id.String()]; ok {
		debug.Log("Index.Lookup", "id %v found in pack %v at %d, length %d",
			id.Str(), p.packID.Str(), p.offset, p.length)
		return p.packID, p.tpe, p.offset, p.length, nil
	}

	debug.Log("Index.Lookup", "id %v not found", id.Str())
	return nil, pack.Data, 0, 0, errors.New("id not found")
}

// Has returns true iff the id is listed in the index.
func (idx *Index) Has(id backend.ID) bool {
	_, _, _, _, err := idx.Lookup(id)
	if err == nil {
		return true
	}

	return false
}

// Merge loads all items from other into idx.
func (idx *Index) Merge(other *Index) {
	debug.Log("Index.Merge", "Merge index with %p", other)
	idx.m.Lock()
	defer idx.m.Unlock()

	for k, v := range other.pack {
		if _, ok := idx.pack[k]; ok {
			debug.Log("Index.Merge", "index already has key %v, updating", k[:8])
		}

		idx.pack[k] = v
	}
	debug.Log("Index.Merge", "done merging index")
}

// Each returns a channel that yields all blobs known to the index. If done is
// closed, the background goroutine terminates. This blocks any modification of
// the index.
func (idx *Index) Each(done chan struct{}) <-chan pack.Blob {
	idx.m.Lock()

	ch := make(chan pack.Blob)

	go func() {
		defer idx.m.Unlock()
		defer func() {
			close(ch)
		}()

		for ids, blob := range idx.pack {
			id, err := backend.ParseID(ids)
			if err != nil {
				// ignore invalid IDs
				continue
			}

			select {
			case <-done:
				return
			case ch <- pack.Blob{
				ID:     id,
				Offset: blob.offset,
				Type:   blob.tpe,
				Length: uint32(blob.length),
			}:
			}
		}
	}()

	return ch
}

// Count returns the number of blobs of type t in the index.
func (idx *Index) Count(t pack.BlobType) (n uint) {
	debug.Log("Index.Count", "counting blobs of type %v", t)
	idx.m.Lock()
	defer idx.m.Unlock()

	for id, blob := range idx.pack {
		if blob.tpe == t {
			n++
			debug.Log("Index.Count", "  blob %v counted: %v", id[:8], blob)
		}
	}

	return
}

type packJSON struct {
	ID    string     `json:"id"`
	Blobs []blobJSON `json:"blobs"`
}

type blobJSON struct {
	ID     string        `json:"id"`
	Type   pack.BlobType `json:"type"`
	Offset uint          `json:"offset"`
	Length uint          `json:"length"`
}

// generatePackList returns a list of packs containing only the index entries
// that selsectFn returned true for. If selectFn is nil, the list contains all
// blobs in the index.
func (idx *Index) generatePackList(selectFn func(indexEntry) bool) ([]*packJSON, error) {
	list := []*packJSON{}
	packs := make(map[string]*packJSON)

	for id, blob := range idx.pack {
		if selectFn != nil && !selectFn(blob) {
			continue
		}

		debug.Log("Index.generatePackList", "handle blob %q", id[:8])

		if blob.packID == nil {
			debug.Log("Index.generatePackList", "blob %q has no packID! (type %v, offset %v, length %v)",
				id[:8], blob.tpe, blob.offset, blob.length)
			return nil, fmt.Errorf("unable to serialize index: pack for blob %v hasn't been written yet", id)
		}

		// see if pack is already in map
		p, ok := packs[blob.packID.String()]
		if !ok {
			// else create new pack
			p = &packJSON{ID: blob.packID.String()}

			// and append it to the list and map
			list = append(list, p)
			packs[p.ID] = p
		}

		// add blob
		p.Blobs = append(p.Blobs, blobJSON{
			ID:     id,
			Type:   blob.tpe,
			Offset: blob.offset,
			Length: blob.length,
		})
	}

	debug.Log("Index.generatePackList", "done")

	return list, nil
}

// encode writes the JSON serialization of the index filtered by selectFn to enc.
func (idx *Index) encode(w io.Writer, selectFn func(indexEntry) bool) error {
	list, err := idx.generatePackList(func(entry indexEntry) bool {
		return !entry.old
	})
	if err != nil {
		return err
	}

	debug.Log("Index.Encode", "done")

	enc := json.NewEncoder(w)
	return enc.Encode(list)
}

// Encode writes the JSON serialization of the index to the writer w. This
// serialization only contains new blobs added via idx.Store(), not old ones
// introduced via DecodeIndex().
func (idx *Index) Encode(w io.Writer) error {
	debug.Log("Index.Encode", "encoding index")
	idx.m.Lock()
	defer idx.m.Unlock()

	return idx.encode(w, func(e indexEntry) bool { return !e.old })
}

// Dump writes the pretty-printed JSON representation of the index to w.
func (idx *Index) Dump(w io.Writer) error {
	debug.Log("Index.Dump", "dumping index")
	idx.m.Lock()
	defer idx.m.Unlock()

	list, err := idx.generatePackList(nil)
	if err != nil {
		return err
	}

	buf, err := json.MarshalIndent(list, "", "  ")
	if err != nil {
		return err
	}

	_, err = w.Write(append(buf, '\n'))
	if err != nil {
		return err
	}

	debug.Log("Index.Dump", "done")

	return nil
}

// DecodeIndex loads and unserializes an index from rd.
func DecodeIndex(rd io.Reader) (*Index, error) {
	debug.Log("Index.DecodeIndex", "Start decoding index")
	list := []*packJSON{}

	dec := json.NewDecoder(rd)
	err := dec.Decode(&list)
	if err != nil {
		return nil, err
	}

	idx := NewIndex()
	for _, pack := range list {
		packID, err := backend.ParseID(pack.ID)
		if err != nil {
			debug.Log("Index.DecodeIndex", "error parsing pack ID %q: %v", pack.ID, err)
			return nil, err
		}

		for _, blob := range pack.Blobs {
			blobID, err := backend.ParseID(blob.ID)
			if err != nil {
				debug.Log("Index.DecodeIndex", "error parsing blob ID %q: %v", blob.ID, err)
				return nil, err
			}

			idx.store(blob.Type, blobID, packID, blob.Offset, blob.Length, true)
		}
	}

	debug.Log("Index.DecodeIndex", "done")
	return idx, err
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"

	"code.google.com/p/gcfg"
	"github.com/rzh/hammer.mongo/hammer"
)

var (
	config    string
	debug     bool
	gen       string
	initdb    bool
	max       bool
	monitor   int64
	profile   string
	quiet     bool
	rps       int64
	run       string // name of the test run
	server    string
	thread    int
	total     int64
	totaltime int64
	warmup    int64
	worker    int
)

type HammerConfig struct {
	Config struct {
		Debug     bool
		Initdb    bool
		Max       bool
		Monitor   int64
		Profile   string
		Quiet     bool
		Rps       int64
		Run       string
		Server    string
		Thread    int
		Total     int64
		Totaltime int64
		Warmup    int64
		Worker    int
	}
}

func parseConfig(cfgFile string) {
	var cfg HammerConfig

	// set up default value first
	cfg.Config.Debug = false
	cfg.Config.Initdb = false
	cfg.Config.Max = false
	cfg.Config.Monitor = 1
	cfg.Config.Profile = ""
	cfg.Config.Quiet = false
	cfg.Config.Rps = 500
	cfg.Config.Server = "localhost:27017"
	cfg.Config.Thread = 0
	cfg.Config.Total = 0
	cfg.Config.Totaltime = 0
	cfg.Config.Warmup = 0
	cfg.Config.Worker = 10

	// to parse config file
	err := gcfg.ReadFileInto(&cfg, cfgFile)
	if err != nil {
		log.Fatalf("Failed to parse configure file %s with error: %s", cfgFile, err)
	}

	// set config flags
	debug = cfg.Config.Debug
	initdb = cfg.Config.Initdb
	max = cfg.Config.Max
	monitor = cfg.Config.Monitor
	profile = cfg.Config.Profile
	quiet = cfg.Config.Quiet
	rps = cfg.Config.Rps
	run = cfg.Config.Run
	server = cfg.Config.Server
	thread = cfg.Config.Thread
	total = cfg.Config.Total
	totaltime = cfg.Config.Totaltime
	warmup = cfg.Config.Warmup
	worker = cfg.Config.Worker

	log.Printf("Configuration file %s parsed, runn with following configurations\n")

	b, err := json.MarshalIndent(cfg, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}

	os.Stdout.Write(b)
}

func init() {
	flag.BoolVar(&debug, "debug", false, "debug flag (true|false)")
	flag.BoolVar(&initdb, "initdb", false, "Drop DB before start testing")
	flag.BoolVar(&max, "max", false, "To find out Max")
	flag.Int64Var(&monitor, "monitor", 1, "Monitor interval")
	flag.StringVar(&profile, "profile", "", "to specify a traffic profile, all UPPERCASE")
	flag.BoolVar(&quiet, "quiet", false, "To silent monitor output")
	flag.Int64Var(&rps, "rps", 500, "Set Request Per Second, 0 to find Max possible")
	flag.StringVar(&run, "run", "", "To specify run id, used for archive test report")
	flag.StringVar(&server, "server", "localhost:27017", "Define server to be tested, default to localhost:27017")
	flag.IntVar(&thread, "thread", 0, "Number of system thread to be used")
	flag.Int64Var(&total, "total", 0, "Total request to be sent, default to unlimited (0)")
	flag.Int64Var(&warmup, "warmup", 0, "To set how long (seconds) for warmup DB")
	flag.IntVar(&worker, "worker", 10, "Number of workers, every worker will have two connections to mongodb")

	flag.Int64Var(&totaltime, "totaltime", 0, "To set how long (seconds) to run the test")
	flag.StringVar(&config, "config", "", "To use config file")
	// flag.StringVar(&gen,       "gen",       "",                "To generate a sample configuration file")
}

func main() {
	flag.Parse()
	if config != "" {
		parseConfig(config)
	}

	NCPU := runtime.NumCPU()

	if thread != 0 {
		runtime.GOMAXPROCS(thread)
	} else {
		runtime.GOMAXPROCS(NCPU + 2) // two more thread than CPU core, need more research on best value TODO
	}

	hammer.Init(worker, monitor, server, initdb, profile, total, warmup, totaltime, quiet, run)

	if max {
		hammer.Start(0)
	} else {
		hammer.Start(rps)
	}

	var input string
	for true {
		fmt.Scanln(&input)

		// todo: make this response to input, to add some CLI options
	}
}
<commit_msg>add commandline options<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"runtime"

	"code.google.com/p/gcfg"
	"github.com/rzh/hammer.mongo/hammer"
)

var (
	config    string
	debug     bool
	gen       string
	initdb    bool
	max       bool
	monitor   int64
	profile   string
	quiet     bool
	rps       int64
	run       string // name of the test run
	server    string
	thread    int
	total     int64
	totaltime int64
	warmup    int64
	worker    int
	ssl bool
	sslca string // PEM file for the CA
	sslkey string // PEM key file for the client
)

type HammerConfig struct {
	Config struct {
		Debug     bool
		Initdb    bool
		Max       bool
		Monitor   int64
		Profile   string
		Quiet     bool
		Rps       int64
		Run       string
		Server    string
		Thread    int
		Total     int64
		Totaltime int64
		Warmup    int64
		Worker    int
	}
}

func parseConfig(cfgFile string) {
	var cfg HammerConfig

	// set up default value first
	cfg.Config.Debug = false
	cfg.Config.Initdb = false
	cfg.Config.Max = false
	cfg.Config.Monitor = 1
	cfg.Config.Profile = ""
	cfg.Config.Quiet = false
	cfg.Config.Rps = 500
	cfg.Config.Server = "localhost:27017"
	cfg.Config.Thread = 0
	cfg.Config.Total = 0
	cfg.Config.Totaltime = 0
	cfg.Config.Warmup = 0
	cfg.Config.Worker = 10

	// to parse config file
	err := gcfg.ReadFileInto(&cfg, cfgFile)
	if err != nil {
		log.Fatalf("Failed to parse configure file %s with error: %s", cfgFile, err)
	}

	// set config flags
	debug = cfg.Config.Debug
	initdb = cfg.Config.Initdb
	max = cfg.Config.Max
	monitor = cfg.Config.Monitor
	profile = cfg.Config.Profile
	quiet = cfg.Config.Quiet
	rps = cfg.Config.Rps
	run = cfg.Config.Run
	server = cfg.Config.Server
	thread = cfg.Config.Thread
	total = cfg.Config.Total
	totaltime = cfg.Config.Totaltime
	warmup = cfg.Config.Warmup
	worker = cfg.Config.Worker

	log.Printf("Configuration file %s parsed, runn with following configurations\n")

	b, err := json.MarshalIndent(cfg, "", "  ")
	if err != nil {
		fmt.Println("error:", err)
	}

	os.Stdout.Write(b)
}

func init() {
	flag.BoolVar(&debug, "debug", false, "debug flag (true|false)")
	flag.BoolVar(&initdb, "initdb", false, "Drop DB before start testing")
	flag.BoolVar(&max, "max", false, "To find out Max")
	flag.Int64Var(&monitor, "monitor", 1, "Monitor interval")
	flag.StringVar(&profile, "profile", "", "to specify a traffic profile, all UPPERCASE")
	flag.BoolVar(&quiet, "quiet", false, "To silent monitor output")
	flag.Int64Var(&rps, "rps", 500, "Set Request Per Second, 0 to find Max possible")
	flag.StringVar(&run, "run", "", "To specify run id, used for archive test report")
	flag.StringVar(&server, "server", "localhost:27017", "Define server to be tested, default to localhost:27017")
	flag.IntVar(&thread, "thread", 0, "Number of system thread to be used")
	flag.Int64Var(&total, "total", 0, "Total request to be sent, default to unlimited (0)")
	flag.Int64Var(&warmup, "warmup", 0, "To set how long (seconds) for warmup DB")
	flag.IntVar(&worker, "worker", 10, "Number of workers, every worker will have two connections to mongodb")
	flag.Int64Var(&totaltime, "totaltime", 0, "To set how long (seconds) to run the test")
	flag.StringVar(&config, "config", "", "To use config file")
	flag.BoolVar(&ssl, "ssl", false, "enable SSL")
	flag.StringVar(&sslca, "sslCAFile", "", "PEM file for CA")
	flag.StringVar(&sslkey, "sslPEMKeyFile", "", "Key file for client")
	// flag.StringVar(&gen,       "gen",       "",                "To generate a sample configuration file")
}

func main() {
	flag.Parse()
	if config != "" {
		parseConfig(config)
	}

	NCPU := runtime.NumCPU()

	if thread != 0 {
		runtime.GOMAXPROCS(thread)
	} else {
		runtime.GOMAXPROCS(NCPU + 2) // two more thread than CPU core, need more research on best value TODO
	}

	hammer.Init(worker, monitor, server, initdb, profile, total, warmup, totaltime, quiet, run)

	if max {
		hammer.Start(0)
	} else {
		hammer.Start(rps)
	}

	var input string
	for true {
		fmt.Scanln(&input)

		// todo: make this response to input, to add some CLI options
	}
}
<|endoftext|>


<commit_before>package acme

import (
	"os"
	"testing"
	"time"

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

var (
	dnsimpleLiveTest bool
	dnsimpleEmail    string
	dnsimpleAPIKey   string
	dnsimpleDomain   string
)

func init() {
	dnsimpleEmail = os.Getenv("DNSIMPLE_EMAIL")
	dnsimpleAPIKey = os.Getenv("DNSIMPLE_API_KEY")
	dnsimpleDomain = os.Getenv("DNSIMPLE_DOMAIN")
	if len(dnsimpleEmail) > 0 && len(dnsimpleAPIKey) > 0 && len(dnsimpleDomain) > 0 {
		dnsimpleLiveTest = true
	}
}

func restoreDNSimpleEnv() {
	os.Setenv("DNSIMPLE_EMAIL", dnsimpleEmail)
	os.Setenv("DNSIMPLE_API_KEY", dnsimpleAPIKey)
}

func TestNewDNSProviderDNSimpleValid(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "")
	os.Setenv("DNSIMPLE_API_KEY", "")
	_, err := NewDNSProviderDNSimple("example@example.com", "123")
	assert.NoError(t, err)
	restoreDNSimpleEnv()
}
func TestNewDNSProviderDNSimpleValidEnv(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "example@example.com")
	os.Setenv("DNSIMPLE_API_KEY", "123")
	_, err := NewDNSProviderDNSimple("", "")
	assert.NoError(t, err)
	restoreDNSimpleEnv()
}

func TestNewDNSProviderDNSimpleMissingCredErr(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "")
	os.Setenv("DNSIMPLE_API_KEY", "")
	_, err := NewDNSProviderDNSimple("", "")
	assert.EqualError(t, err, "DNSimple credentials missing")
	restoreDNSimpleEnv()
}

func TestLiveDNSimplePresent(t *testing.T) {
	if !dnsimpleLiveTest {
		t.Skip("skipping live test")
	}

	provider, err := NewDNSProviderDNSimple(dnsimpleEmail, dnsimpleAPIKey)
	assert.NoError(t, err)

	err = provider.Present(dnsimpleDomain, "", "123d==")
	assert.NoError(t, err)
}

func TestLiveDNSimpleCleanUp(t *testing.T) {
	if !dnsimpleLiveTest {
		t.Skip("skipping live test")
	}

	time.Sleep(time.Second * 1)

	provider, err := NewDNSProviderDNSimple(cflareEmail, cflareAPIKey)
	assert.NoError(t, err)

	err = provider.CleanUp(dnsimpleDomain, "", "123d==")
	assert.NoError(t, err)
}
<commit_msg>Fix wrong variables being used in DNSimple test.<commit_after>package acme

import (
	"os"
	"testing"
	"time"

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

var (
	dnsimpleLiveTest bool
	dnsimpleEmail    string
	dnsimpleAPIKey   string
	dnsimpleDomain   string
)

func init() {
	dnsimpleEmail = os.Getenv("DNSIMPLE_EMAIL")
	dnsimpleAPIKey = os.Getenv("DNSIMPLE_API_KEY")
	dnsimpleDomain = os.Getenv("DNSIMPLE_DOMAIN")
	if len(dnsimpleEmail) > 0 && len(dnsimpleAPIKey) > 0 && len(dnsimpleDomain) > 0 {
		dnsimpleLiveTest = true
	}
}

func restoreDNSimpleEnv() {
	os.Setenv("DNSIMPLE_EMAIL", dnsimpleEmail)
	os.Setenv("DNSIMPLE_API_KEY", dnsimpleAPIKey)
}

func TestNewDNSProviderDNSimpleValid(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "")
	os.Setenv("DNSIMPLE_API_KEY", "")
	_, err := NewDNSProviderDNSimple("example@example.com", "123")
	assert.NoError(t, err)
	restoreDNSimpleEnv()
}
func TestNewDNSProviderDNSimpleValidEnv(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "example@example.com")
	os.Setenv("DNSIMPLE_API_KEY", "123")
	_, err := NewDNSProviderDNSimple("", "")
	assert.NoError(t, err)
	restoreDNSimpleEnv()
}

func TestNewDNSProviderDNSimpleMissingCredErr(t *testing.T) {
	os.Setenv("DNSIMPLE_EMAIL", "")
	os.Setenv("DNSIMPLE_API_KEY", "")
	_, err := NewDNSProviderDNSimple("", "")
	assert.EqualError(t, err, "DNSimple credentials missing")
	restoreDNSimpleEnv()
}

func TestLiveDNSimplePresent(t *testing.T) {
	if !dnsimpleLiveTest {
		t.Skip("skipping live test")
	}

	provider, err := NewDNSProviderDNSimple(dnsimpleEmail, dnsimpleAPIKey)
	assert.NoError(t, err)

	err = provider.Present(dnsimpleDomain, "", "123d==")
	assert.NoError(t, err)
}

func TestLiveDNSimpleCleanUp(t *testing.T) {
	if !dnsimpleLiveTest {
		t.Skip("skipping live test")
	}

	time.Sleep(time.Second * 1)

	provider, err := NewDNSProviderDNSimple(dnsimpleEmail, dnsimpleAPIKey)
	assert.NoError(t, err)

	err = provider.CleanUp(dnsimpleDomain, "", "123d==")
	assert.NoError(t, err)
}
<|endoftext|>


<commit_before>package handler_test

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"

	"context"

	"github.com/graphql-go/graphql"
	"github.com/graphql-go/graphql/testutil"
	"github.com/graphql-go/handler"
)

func decodeResponse(t *testing.T, recorder *httptest.ResponseRecorder) *graphql.Result {
	// clone request body reader so that we can have a nicer error message
	bodyString := ""
	var target graphql.Result
	if b, err := ioutil.ReadAll(recorder.Body); err == nil {
		bodyString = string(b)
	}
	readerClone := strings.NewReader(bodyString)

	decoder := json.NewDecoder(readerClone)
	err := decoder.Decode(&target)
	if err != nil {
		t.Fatalf("DecodeResponseToType(): %v \n%v", err.Error(), bodyString)
	}
	return &target
}
func executeTest(t *testing.T, h *handler.Handler, req *http.Request) (*graphql.Result, *httptest.ResponseRecorder) {
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	result := decodeResponse(t, resp)
	return result, resp
}

func TestContextPropagated(t *testing.T) {
	myNameQuery := graphql.NewObject(graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"name": &graphql.Field{
				Name: "name",
				Type: graphql.String,
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					return p.Context.Value("name"), nil
				},
			},
		},
	})
	myNameSchema, err := graphql.NewSchema(graphql.SchemaConfig{
		Query: myNameQuery,
	})
	if err != nil {
		t.Fatal(err)
	}

	expected := &graphql.Result{
		Data: map[string]interface{}{
			"name": "context-data",
		},
	}
	queryString := `query={name}`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &myNameSchema,
		Pretty: true,
	})

	ctx := context.WithValue(context.Background(), "name", "context-data")
	resp := httptest.NewRecorder()
	h.ContextHandler(ctx, resp, req)
	result := decodeResponse(t, resp)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}

func TestHandler_BasicQuery_Pretty(t *testing.T) {
	expected := &graphql.Result{
		Data: map[string]interface{}{
			"hero": map[string]interface{}{
				"name": "R2-D2",
			},
		},
	}
	queryString := `query=query HeroNameQuery { hero { name } }&operationName=HeroNameQuery`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	callbackCalled := false
	h := handler.New(&handler.Config{
		Schema: &testutil.StarWarsSchema,
		Pretty: true,
		ResultCallbackFn: func(ctx context.Context, params *graphql.Params, result *graphql.Result, responseBody []byte) {
			callbackCalled = true
			if params.OperationName != "HeroNameQuery" {
				t.Fatalf("OperationName passed to callback was not HeroNameQuery: %v", params.OperationName)
			}

			if result.HasErrors() {
				t.Fatalf("unexpected graphql result errors")
			}
		},
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
	if !callbackCalled {
		t.Fatalf("ResultCallbackFn was not called when it should have been")
	}
}

func TestHandler_BasicQuery_Ugly(t *testing.T) {
	expected := &graphql.Result{
		Data: map[string]interface{}{
			"hero": map[string]interface{}{
				"name": "R2-D2",
			},
		},
	}
	queryString := `query=query HeroNameQuery { hero { name } }`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &testutil.StarWarsSchema,
		Pretty: false,
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}

func TestHandler_Params_NilParams(t *testing.T) {
	defer func() {
		if r := recover(); r != nil {
			if str, ok := r.(string); ok {
				if str != "undefined GraphQL schema" {
					t.Fatalf("unexpected error, got %v", r)
				}
				// test passed
				return
			}
			t.Fatalf("unexpected error, got %v", r)

		}
		t.Fatalf("expected to panic, did not panic")
	}()
	_ = handler.New(nil)

}

func TestHandler_BasicQuery_WithRootObjFn(t *testing.T) {
	myNameQuery := graphql.NewObject(graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"name": &graphql.Field{
				Name: "name",
				Type: graphql.String,
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					rv := p.Info.RootValue.(map[string]interface{})
					return rv["rootValue"], nil
				},
			},
		},
	})
	myNameSchema, err := graphql.NewSchema(graphql.SchemaConfig{
		Query: myNameQuery,
	})
	if err != nil {
		t.Fatal(err)
	}

	expected := &graphql.Result{
		Data: map[string]interface{}{
			"name": "foo",
		},
	}
	queryString := `query={name}`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &myNameSchema,
		Pretty: true,
		RootObjectFn: func(ctx context.Context, r *http.Request) map[string]interface{} {
			return map[string]interface{}{"rootValue": "foo"}
		},
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}
<commit_msg>handler_test: adds FormatErrorFn tests<commit_after>package handler_test

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"reflect"
	"strings"
	"testing"

	"context"

	"github.com/graphql-go/graphql"
	"github.com/graphql-go/graphql/gqlerrors"
	"github.com/graphql-go/graphql/language/location"
	"github.com/graphql-go/graphql/testutil"
	"github.com/graphql-go/handler"
)

func decodeResponse(t *testing.T, recorder *httptest.ResponseRecorder) *graphql.Result {
	// clone request body reader so that we can have a nicer error message
	bodyString := ""
	var target graphql.Result
	if b, err := ioutil.ReadAll(recorder.Body); err == nil {
		bodyString = string(b)
	}
	readerClone := strings.NewReader(bodyString)

	decoder := json.NewDecoder(readerClone)
	err := decoder.Decode(&target)
	if err != nil {
		t.Fatalf("DecodeResponseToType(): %v \n%v", err.Error(), bodyString)
	}
	return &target
}
func executeTest(t *testing.T, h *handler.Handler, req *http.Request) (*graphql.Result, *httptest.ResponseRecorder) {
	resp := httptest.NewRecorder()
	h.ServeHTTP(resp, req)
	result := decodeResponse(t, resp)
	return result, resp
}

func TestContextPropagated(t *testing.T) {
	myNameQuery := graphql.NewObject(graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"name": &graphql.Field{
				Name: "name",
				Type: graphql.String,
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					return p.Context.Value("name"), nil
				},
			},
		},
	})
	myNameSchema, err := graphql.NewSchema(graphql.SchemaConfig{
		Query: myNameQuery,
	})
	if err != nil {
		t.Fatal(err)
	}

	expected := &graphql.Result{
		Data: map[string]interface{}{
			"name": "context-data",
		},
	}
	queryString := `query={name}`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &myNameSchema,
		Pretty: true,
	})

	ctx := context.WithValue(context.Background(), "name", "context-data")
	resp := httptest.NewRecorder()
	h.ContextHandler(ctx, resp, req)
	result := decodeResponse(t, resp)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}

func TestHandler_BasicQuery_Pretty(t *testing.T) {
	expected := &graphql.Result{
		Data: map[string]interface{}{
			"hero": map[string]interface{}{
				"name": "R2-D2",
			},
		},
	}
	queryString := `query=query HeroNameQuery { hero { name } }&operationName=HeroNameQuery`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	callbackCalled := false
	h := handler.New(&handler.Config{
		Schema: &testutil.StarWarsSchema,
		Pretty: true,
		ResultCallbackFn: func(ctx context.Context, params *graphql.Params, result *graphql.Result, responseBody []byte) {
			callbackCalled = true
			if params.OperationName != "HeroNameQuery" {
				t.Fatalf("OperationName passed to callback was not HeroNameQuery: %v", params.OperationName)
			}

			if result.HasErrors() {
				t.Fatalf("unexpected graphql result errors")
			}
		},
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
	if !callbackCalled {
		t.Fatalf("ResultCallbackFn was not called when it should have been")
	}
}

func TestHandler_BasicQuery_Ugly(t *testing.T) {
	expected := &graphql.Result{
		Data: map[string]interface{}{
			"hero": map[string]interface{}{
				"name": "R2-D2",
			},
		},
	}
	queryString := `query=query HeroNameQuery { hero { name } }`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &testutil.StarWarsSchema,
		Pretty: false,
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}

func TestHandler_Params_NilParams(t *testing.T) {
	defer func() {
		if r := recover(); r != nil {
			if str, ok := r.(string); ok {
				if str != "undefined GraphQL schema" {
					t.Fatalf("unexpected error, got %v", r)
				}
				// test passed
				return
			}
			t.Fatalf("unexpected error, got %v", r)

		}
		t.Fatalf("expected to panic, did not panic")
	}()
	_ = handler.New(nil)

}

func TestHandler_BasicQuery_WithRootObjFn(t *testing.T) {
	myNameQuery := graphql.NewObject(graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"name": &graphql.Field{
				Name: "name",
				Type: graphql.String,
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					rv := p.Info.RootValue.(map[string]interface{})
					return rv["rootValue"], nil
				},
			},
		},
	})
	myNameSchema, err := graphql.NewSchema(graphql.SchemaConfig{
		Query: myNameQuery,
	})
	if err != nil {
		t.Fatal(err)
	}

	expected := &graphql.Result{
		Data: map[string]interface{}{
			"name": "foo",
		},
	}
	queryString := `query={name}`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	h := handler.New(&handler.Config{
		Schema: &myNameSchema,
		Pretty: true,
		RootObjectFn: func(ctx context.Context, r *http.Request) map[string]interface{} {
			return map[string]interface{}{"rootValue": "foo"}
		},
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}

type customError struct {
	error
}

func (e customError) Error() string {
	return e.error.Error()
}

func TestHandler_BasicQuery_WithFormatErrorFn(t *testing.T) {
	resolverError := customError{error: errors.New("resolver error")}
	myNameQuery := graphql.NewObject(graphql.ObjectConfig{
		Name: "Query",
		Fields: graphql.Fields{
			"name": &graphql.Field{
				Name: "name",
				Type: graphql.String,
				Resolve: func(p graphql.ResolveParams) (interface{}, error) {
					return nil, resolverError
				},
			},
		},
	})
	myNameSchema, err := graphql.NewSchema(graphql.SchemaConfig{
		Query: myNameQuery,
	})
	if err != nil {
		t.Fatal(err)
	}

	customFormattedError := gqlerrors.FormattedError{
		Message: resolverError.Error(),
		Locations: []location.SourceLocation{
			location.SourceLocation{
				Line:   1,
				Column: 2,
			},
		},
		Path: []interface{}{"name"},
		Extensions: map[string]interface{}{
			"fromFormatFn": "FROM_FORMAT_FN",
		},
	}

	expected := &graphql.Result{
		Data: map[string]interface{}{
			"name": nil,
		},
		Errors: []gqlerrors.FormattedError{customFormattedError},
	}

	queryString := `query={name}`
	req, _ := http.NewRequest("GET", fmt.Sprintf("/graphql?%v", queryString), nil)

	formatErrorFnCalled := false
	h := handler.New(&handler.Config{
		Schema: &myNameSchema,
		Pretty: true,
		FormatErrorFn: func(err gqlerrors.FormattedError) gqlerrors.FormattedError {
			formatErrorFnCalled = true
			originalError := err.OriginalError()
			switch errType := originalError.(type) {
			case customError:
			default:
				t.Fatalf("unexpected error type: %v", reflect.TypeOf(errType))
			}
			return gqlerrors.FormattedError{
				Message:   err.Message,
				Locations: err.Locations,
				Path:      err.Path,
				Extensions: map[string]interface{}{
					"fromFormatFn": "FROM_FORMAT_FN",
				},
			}
		},
	})
	result, resp := executeTest(t, h, req)
	if resp.Code != http.StatusOK {
		t.Fatalf("unexpected server response %v", resp.Code)
	}
	if !formatErrorFnCalled {
		t.Fatalf("FormatErrorFn was not called when it should have been")
	}
	if !reflect.DeepEqual(result, expected) {
		t.Fatalf("wrong result, graphql result diff: %v", testutil.Diff(expected, result))
	}
}
<|endoftext|>