<commit_before>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package agent

import (
	"errors"
	"io"
	"net"
	"sync"

	"code.google.com/p/go.crypto/ssh"
)

// RequestAgentForwarding sets up agent forwarding for the session.
// SetupForwardKeyring or SetupForwardAgent should be called to route
// the authentication requests.
func RequestAgentForwarding(session *ssh.Session) error {
	ok, err := session.SendRequest("auth-agent-req@openssh.com", true, nil)
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("forwarding request denied")
	}
	return nil
}

// ForwardToAgent routes authentication requests to the given keyring.
func ForwardToAgent(client *ssh.Client, keyring Agent) error {
	channels := client.HandleChannelOpen(channelType)
	if channels == nil {
		return errors.New("agent: already have handler for " + channelType)
	}

	go func() {
		for ch := range channels {
			channel, reqs, err := ch.Accept()
			if err != nil {
				continue
			}
			go ssh.DiscardRequests(reqs)
			go func() {
				ServeAgent(keyring, channel)
				channel.Close()
			}()
		}
	}()
	return nil
}

const channelType = "auth-agent@openssh.com"

// ForwardToRemote routes authentication requests to the ssh-agent
// process serving on the given unix socket.
func ForwardToRemote(client *ssh.Client, addr string) error {
	channels := client.HandleChannelOpen(channelType)
	if channels == nil {
		return errors.New("agent: already have handler for " + channelType)
	}
	conn, err := net.Dial("unix", addr)
	if err != nil {
		return err
	}
	conn.Close()

	go func() {
		for ch := range channels {
			channel, reqs, err := ch.Accept()
			if err != nil {
				continue
			}
			go ssh.DiscardRequests(reqs)
			go forwardUnixSocket(channel, addr)
		}
	}()
	return nil
}

func forwardUnixSocket(channel ssh.Channel, addr string) {
	conn, err := net.Dial("unix", addr)
	if err != nil {
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(conn, channel)
		conn.(*net.UnixConn).CloseWrite()
		wg.Done()
	}()
	go func() {
		io.Copy(channel, conn)
		channel.CloseWrite()
		wg.Done()
	}()

	wg.Wait()
	conn.Close()
	channel.Close()
}
<commit_msg>ssh/forward: Fix reference to the functions to call to route authentication requests.<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package agent

import (
	"errors"
	"io"
	"net"
	"sync"

	"code.google.com/p/go.crypto/ssh"
)

// RequestAgentForwarding sets up agent forwarding for the session.
// ForwardToAgent or ForwardToRemote should be called to route
// the authentication requests.
func RequestAgentForwarding(session *ssh.Session) error {
	ok, err := session.SendRequest("auth-agent-req@openssh.com", true, nil)
	if err != nil {
		return err
	}
	if !ok {
		return errors.New("forwarding request denied")
	}
	return nil
}

// ForwardToAgent routes authentication requests to the given keyring.
func ForwardToAgent(client *ssh.Client, keyring Agent) error {
	channels := client.HandleChannelOpen(channelType)
	if channels == nil {
		return errors.New("agent: already have handler for " + channelType)
	}

	go func() {
		for ch := range channels {
			channel, reqs, err := ch.Accept()
			if err != nil {
				continue
			}
			go ssh.DiscardRequests(reqs)
			go func() {
				ServeAgent(keyring, channel)
				channel.Close()
			}()
		}
	}()
	return nil
}

const channelType = "auth-agent@openssh.com"

// ForwardToRemote routes authentication requests to the ssh-agent
// process serving on the given unix socket.
func ForwardToRemote(client *ssh.Client, addr string) error {
	channels := client.HandleChannelOpen(channelType)
	if channels == nil {
		return errors.New("agent: already have handler for " + channelType)
	}
	conn, err := net.Dial("unix", addr)
	if err != nil {
		return err
	}
	conn.Close()

	go func() {
		for ch := range channels {
			channel, reqs, err := ch.Accept()
			if err != nil {
				continue
			}
			go ssh.DiscardRequests(reqs)
			go forwardUnixSocket(channel, addr)
		}
	}()
	return nil
}

func forwardUnixSocket(channel ssh.Channel, addr string) {
	conn, err := net.Dial("unix", addr)
	if err != nil {
		return
	}

	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		io.Copy(conn, channel)
		conn.(*net.UnixConn).CloseWrite()
		wg.Done()
	}()
	go func() {
		io.Copy(channel, conn)
		channel.CloseWrite()
		wg.Done()
	}()

	wg.Wait()
	conn.Close()
	channel.Close()
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"net/http"

	"github.com/gilcrest/go-API-template/pkg/api/server/handlers"
	"github.com/gilcrest/go-API-template/pkg/env"

	"github.com/gorilla/mux"
)

func main() {

	// Initializes "environment" struct type
	env, err := env.NewEnv()

	if err != nil {
		log.Fatal(err)
	}

	// create a new mux (multiplex) router
	rtr := mux.NewRouter()

	// API may have multiple versions and the matching may get a bit
	// lengthy, this RouteMatch function helps with organizing that
	rtr = handlers.Dispatch(env, rtr)

	// handle all requests with the Gorilla router by adding
	// rtr to the DefaultServeMux
	// LogRequest middleware will log request as well
	http.Handle("/", rtr)

	if err := http.ListenAndServe("127.0.0.1:8080", nil); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Added call to NewSubrouter<commit_after>package main

import (
	"log"
	"net/http"

	"github.com/gilcrest/go-API-template/pkg/api/server/handlers"
	"github.com/gilcrest/go-API-template/pkg/env"

	"github.com/gorilla/mux"
)

func main() {

	// Initializes "environment" struct type
	env, err := env.NewEnv()

	if err != nil {
		log.Fatal(err)
	}

	// create a new mux (multiplex) router
	rtr := mux.NewRouter()

	// send Router through subRouter function to add any standard
	// Subroutes you may want for your APIs
	r := handlers.NewSubrouter(rtr)

	// API may have multiple versions and the matching may get a bit
	// lengthy, this RouteMatch function helps with organizing that
	r = handlers.Dispatch(env, r)

	// handle all requests with the Gorilla router by adding
	// rtr to the DefaultServeMux
	// LogRequest middleware will log request as well
	http.Handle("/", r)

	if err := http.ListenAndServe("127.0.0.1:8080", nil); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package run_test

import (
	"os"
	"testing"

	"github.com/BurntSushi/toml"
	"github.com/influxdata/kapacitor/cmd/kapacitord/run"
)

// Ensure the configuration can be parsed.
func TestConfig_Parse(t *testing.T) {
	// Parse configuration.
	var c run.Config
	if _, err := toml.Decode(`
[replay]
dir = "/tmp/replay"

[task]
dir = "/tmp/task"
`, &c); err != nil {
		t.Fatal(err)
	}

	// Validate configuration.
	if c.Replay.Dir != "/tmp/replay" {
		t.Fatalf("unexpected replay dir: %s", c.Replay.Dir)
	} else if c.Task.Dir != "/tmp/task" {
		t.Fatalf("unexpected task dir: %s", c.Task.Dir)
	}
}

// Ensure the configuration can be parsed.
func TestConfig_Parse_EnvOverride(t *testing.T) {
	// Parse configuration.
	var c run.Config
	if _, err := toml.Decode(`
[replay]
dir = "/tmp/replay"

[task]
dir = "/tmp/task"
`, &c); err != nil {
		t.Fatal(err)
	}

	if err := os.Setenv("KAPACITOR_REPLAY_DIR", "/var/lib/kapacitor/replay"); err != nil {
		t.Fatalf("failed to set env var: %v", err)
	}

	if err := os.Setenv("KAPACITOR_TASK_DIR", "/var/lib/kapacitor/task"); err != nil {
		t.Fatalf("failed to set env var: %v", err)
	}

	if err := c.ApplyEnvOverrides(); err != nil {
		t.Fatalf("failed to apply env overrides: %v", err)
	}

	// Validate configuration.
	if c.Replay.Dir != "/var/lib/kapacitor/replay" {
		t.Fatalf("unexpected replay dir: %s", c.Replay.Dir)
	} else if c.Task.Dir != "/var/lib/kapacitor/task" {
		t.Fatalf("unexpected task dir: %s", c.Task.Dir)
	}
}
<commit_msg>config_test: add test for overriding array/slice elements.<commit_after>package run_test

import (
	"os"
	"testing"

	"github.com/BurntSushi/toml"
	"github.com/influxdata/kapacitor/cmd/kapacitord/run"
)

// Ensure the configuration can be parsed.
func TestConfig_Parse(t *testing.T) {
	// Parse configuration.
	var c run.Config
	if _, err := toml.Decode(`
[replay]
dir = "/tmp/replay"

[task]
dir = "/tmp/task"
`, &c); err != nil {
		t.Fatal(err)
	}

	// Validate configuration.
	if c.Replay.Dir != "/tmp/replay" {
		t.Fatalf("unexpected replay dir: %s", c.Replay.Dir)
	} else if c.Task.Dir != "/tmp/task" {
		t.Fatalf("unexpected task dir: %s", c.Task.Dir)
	}
}

// Ensure the configuration can be parsed.
func TestConfig_Parse_EnvOverride(t *testing.T) {
	// Parse configuration.
	var c run.Config
	if _, err := toml.Decode(`
[influxdb]
urls=["http://localhost:8086"]

[replay]
dir = "/tmp/replay"

[task]
dir = "/tmp/task"
`, &c); err != nil {
		t.Fatal(err)
	}

	if err := os.Setenv("KAPACITOR_REPLAY_DIR", "/var/lib/kapacitor/replay"); err != nil {
		t.Fatalf("failed to set env var: %v", err)
	}

	if err := os.Setenv("KAPACITOR_TASK_DIR", "/var/lib/kapacitor/task"); err != nil {
		t.Fatalf("failed to set env var: %v", err)
	}

	if err := os.Setenv("KAPACITOR_INFLUXDB_URLS_0", "http://localhost:18086"); err != nil {
		t.Fatalf("failed to set env var: %v", err)
	}

	if err := c.ApplyEnvOverrides(); err != nil {
		t.Fatalf("failed to apply env overrides: %v", err)
	}

	// Validate configuration.
	if c.Replay.Dir != "/var/lib/kapacitor/replay" {
		t.Fatalf("unexpected replay dir: %s", c.Replay.Dir)
	} else if c.Task.Dir != "/var/lib/kapacitor/task" {
		t.Fatalf("unexpected task dir: %s", c.Task.Dir)
	} else if c.InfluxDB.URLs[0] != "http://localhost:18086" {
		t.Fatalf("unexpected url 0: %s", c.InfluxDB.URLs[0])
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"bytes"
	"fmt"
	"io"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	kubeadmutil "k8s.io/kubernetes/cmd/kubeadm/app/util"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

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

var (
	completion_long = dedent.Dedent(`
		Output shell completion code for the specified shell (bash or zsh).
		The shell code must be evalutated to provide interactive
		completion of kubeadm commands.  This can be done by sourcing it from
		the .bash_profile.

		Note: this requires the bash-completion framework, which is not installed
		by default on Mac.  This can be installed by using homebrew:

		    $ brew install bash-completion

		Once installed, bash_completion must be evaluated.  This can be done by adding the
		following line to the .bash_profile

		    $ source $(brew --prefix)/etc/bash_completion

		Note for zsh users: [1] zsh completions are only supported in versions of zsh >= 5.2`)

	completion_example = dedent.Dedent(`
		# Install bash completion on a Mac using homebrew
		brew install bash-completion
		printf "\n# Bash completion support\nsource $(brew --prefix)/etc/bash_completion\n" >> $HOME/.bash_profile
		source $HOME/.bash_profile

		# Load the kubeadm completion code for bash into the current shell
		source <(kubeadm completion bash)

		# Write bash completion code to a file and source if from .bash_profile
		kubeadm completion bash > ~/.kube/kubeadm_completion.bash.inc
		printf "\n# Kubeadm shell completion\nsource '$HOME/.kube/kubeadm_completion.bash.inc'\n" >> $HOME/.bash_profile
		source $HOME/.bash_profile

		# Load the kubeadm completion code for zsh[1] into the current shell
		source <(kubeadm completion zsh)`)
)

var (
	completion_shells = map[string]func(out io.Writer, cmd *cobra.Command) error{
		"bash": runCompletionBash,
		"zsh":  runCompletionZsh,
	}
)

func NewCmdCompletion(out io.Writer, boilerPlate string) *cobra.Command {
	shells := []string{}
	for s := range completion_shells {
		shells = append(shells, s)
	}

	cmd := &cobra.Command{
		Use:     "completion SHELL",
		Short:   i18n.T("Output shell completion code for the specified shell (bash or zsh)"),
		Long:    completion_long,
		Example: completion_example,
		Run: func(cmd *cobra.Command, args []string) {
			err := RunCompletion(out, boilerPlate, cmd, args)
			kubeadmutil.CheckErr(err)
		},
		ValidArgs: shells,
	}

	return cmd
}

func RunCompletion(out io.Writer, boilerPlate string, cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("shell not specified.")
	}
	if len(args) > 1 {
		return fmt.Errorf("too many arguments. expected only the shell type.")
	}
	run, found := completion_shells[args[0]]
	if !found {
		return fmt.Errorf("unsupported shell type %q.", args[0])
	}

	if len(boilerPlate) == 0 {
		boilerPlate = defaultBoilerPlate
	}
	if _, err := out.Write([]byte(boilerPlate)); err != nil {
		return err
	}
	return run(out, cmd.Parent())
}

func runCompletionBash(out io.Writer, kubeadm *cobra.Command) error {
	return kubeadm.GenBashCompletion(out)
}

func runCompletionZsh(out io.Writer, kubeadm *cobra.Command) error {
	zsh_initialization := `
__kubeadm_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand

	source "$@"
}

__kubeadm_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift

		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__kubeadm_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}

__kubeadm_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?

	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}

__kubeadm_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}

__kubeadm_declare() {
	if [ "$1" == "-F" ]; then
		whence -w "$@"
	else
		builtin declare "$@"
	fi
}

__kubeadm_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}

__kubeadm_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}

__kubeadm_filedir() {
	local RET OLD_IFS w qw

	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi

	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"

	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"

	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__kubeadm_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}

__kubeadm_quote() {
    if [[ $1 == \'* || $1 == \"* ]]; then
        # Leave out first character
        printf %q "${1:1}"
    else
    	printf %q "$1"
    fi
}

autoload -U +X bashcompinit && bashcompinit

# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi

__kubeadm_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__kubeadm_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__kubeadm_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__kubeadm_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__kubeadm_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__kubeadm_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/__kubeadm_declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__kubeadm_type/g" \
	<<'BASH_COMPLETION_EOF'
`
	out.Write([]byte(zsh_initialization))

	buf := new(bytes.Buffer)
	kubeadm.GenBashCompletion(buf)
	out.Write(buf.Bytes())

	zsh_tail := `
BASH_COMPLETION_EOF
}

__kubeadm_bash_source <(__kubeadm_convert_bash_to_zsh)
`
	out.Write([]byte(zsh_tail))
	return nil
}
<commit_msg>Fix zsh completion for kubeadm<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package cmd

import (
	"bytes"
	"fmt"
	"io"

	"github.com/renstrom/dedent"
	"github.com/spf13/cobra"

	kubeadmutil "k8s.io/kubernetes/cmd/kubeadm/app/util"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
)

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

var (
	completion_long = dedent.Dedent(`
		Output shell completion code for the specified shell (bash or zsh).
		The shell code must be evalutated to provide interactive
		completion of kubeadm commands.  This can be done by sourcing it from
		the .bash_profile.

		Note: this requires the bash-completion framework, which is not installed
		by default on Mac.  This can be installed by using homebrew:

		    $ brew install bash-completion

		Once installed, bash_completion must be evaluated.  This can be done by adding the
		following line to the .bash_profile

		    $ source $(brew --prefix)/etc/bash_completion

		Note for zsh users: [1] zsh completions are only supported in versions of zsh >= 5.2`)

	completion_example = dedent.Dedent(`
		# Install bash completion on a Mac using homebrew
		brew install bash-completion
		printf "\n# Bash completion support\nsource $(brew --prefix)/etc/bash_completion\n" >> $HOME/.bash_profile
		source $HOME/.bash_profile

		# Load the kubeadm completion code for bash into the current shell
		source <(kubeadm completion bash)

		# Write bash completion code to a file and source if from .bash_profile
		kubeadm completion bash > ~/.kube/kubeadm_completion.bash.inc
		printf "\n# Kubeadm shell completion\nsource '$HOME/.kube/kubeadm_completion.bash.inc'\n" >> $HOME/.bash_profile
		source $HOME/.bash_profile

		# Load the kubeadm completion code for zsh[1] into the current shell
		source <(kubeadm completion zsh)`)
)

var (
	completion_shells = map[string]func(out io.Writer, cmd *cobra.Command) error{
		"bash": runCompletionBash,
		"zsh":  runCompletionZsh,
	}
)

func NewCmdCompletion(out io.Writer, boilerPlate string) *cobra.Command {
	shells := []string{}
	for s := range completion_shells {
		shells = append(shells, s)
	}

	cmd := &cobra.Command{
		Use:     "completion SHELL",
		Short:   i18n.T("Output shell completion code for the specified shell (bash or zsh)"),
		Long:    completion_long,
		Example: completion_example,
		Run: func(cmd *cobra.Command, args []string) {
			err := RunCompletion(out, boilerPlate, cmd, args)
			kubeadmutil.CheckErr(err)
		},
		ValidArgs: shells,
	}

	return cmd
}

func RunCompletion(out io.Writer, boilerPlate string, cmd *cobra.Command, args []string) error {
	if len(args) == 0 {
		return fmt.Errorf("shell not specified.")
	}
	if len(args) > 1 {
		return fmt.Errorf("too many arguments. expected only the shell type.")
	}
	run, found := completion_shells[args[0]]
	if !found {
		return fmt.Errorf("unsupported shell type %q.", args[0])
	}

	if len(boilerPlate) == 0 {
		boilerPlate = defaultBoilerPlate
	}
	if _, err := out.Write([]byte(boilerPlate)); err != nil {
		return err
	}
	return run(out, cmd.Parent())
}

func runCompletionBash(out io.Writer, kubeadm *cobra.Command) error {
	return kubeadm.GenBashCompletion(out)
}

func runCompletionZsh(out io.Writer, kubeadm *cobra.Command) error {
	zsh_initialization := `
__kubeadm_bash_source() {
	alias shopt=':'
	alias _expand=_bash_expand
	alias _complete=_bash_comp
	emulate -L sh
	setopt kshglob noshglob braceexpand

	source "$@"
}

__kubeadm_type() {
	# -t is not supported by zsh
	if [ "$1" == "-t" ]; then
		shift

		# fake Bash 4 to disable "complete -o nospace". Instead
		# "compopt +-o nospace" is used in the code to toggle trailing
		# spaces. We don't support that, but leave trailing spaces on
		# all the time
		if [ "$1" = "__kubeadm_compopt" ]; then
			echo builtin
			return 0
		fi
	fi
	type "$@"
}

__kubeadm_compgen() {
	local completions w
	completions=( $(compgen "$@") ) || return $?

	# filter by given word as prefix
	while [[ "$1" = -* && "$1" != -- ]]; do
		shift
		shift
	done
	if [[ "$1" == -- ]]; then
		shift
	fi
	for w in "${completions[@]}"; do
		if [[ "${w}" = "$1"* ]]; then
			echo "${w}"
		fi
	done
}

__kubeadm_compopt() {
	true # don't do anything. Not supported by bashcompinit in zsh
}

__kubeadm_ltrim_colon_completions()
{
	if [[ "$1" == *:* && "$COMP_WORDBREAKS" == *:* ]]; then
		# Remove colon-word prefix from COMPREPLY items
		local colon_word=${1%${1##*:}}
		local i=${#COMPREPLY[*]}
		while [[ $((--i)) -ge 0 ]]; do
			COMPREPLY[$i]=${COMPREPLY[$i]#"$colon_word"}
		done
	fi
}

__kubeadm_get_comp_words_by_ref() {
	cur="${COMP_WORDS[COMP_CWORD]}"
	prev="${COMP_WORDS[${COMP_CWORD}-1]}"
	words=("${COMP_WORDS[@]}")
	cword=("${COMP_CWORD[@]}")
}

__kubeadm_filedir() {
	local RET OLD_IFS w qw

	__debug "_filedir $@ cur=$cur"
	if [[ "$1" = \~* ]]; then
		# somehow does not work. Maybe, zsh does not call this at all
		eval echo "$1"
		return 0
	fi

	OLD_IFS="$IFS"
	IFS=$'\n'
	if [ "$1" = "-d" ]; then
		shift
		RET=( $(compgen -d) )
	else
		RET=( $(compgen -f) )
	fi
	IFS="$OLD_IFS"

	IFS="," __debug "RET=${RET[@]} len=${#RET[@]}"

	for w in ${RET[@]}; do
		if [[ ! "${w}" = "${cur}"* ]]; then
			continue
		fi
		if eval "[[ \"\${w}\" = *.$1 || -d \"\${w}\" ]]"; then
			qw="$(__kubeadm_quote "${w}")"
			if [ -d "${w}" ]; then
				COMPREPLY+=("${qw}/")
			else
				COMPREPLY+=("${qw}")
			fi
		fi
	done
}

__kubeadm_quote() {
    if [[ $1 == \'* || $1 == \"* ]]; then
        # Leave out first character
        printf %q "${1:1}"
    else
    	printf %q "$1"
    fi
}

autoload -U +X bashcompinit && bashcompinit

# use word boundary patterns for BSD or GNU sed
LWORD='[[:<:]]'
RWORD='[[:>:]]'
if sed --help 2>&1 | grep -q GNU; then
	LWORD='\<'
	RWORD='\>'
fi

__kubeadm_convert_bash_to_zsh() {
	sed \
	-e 's/declare -F/whence -w/' \
	-e 's/local \([a-zA-Z0-9_]*\)=/local \1; \1=/' \
	-e 's/flags+=("\(--.*\)=")/flags+=("\1"); two_word_flags+=("\1")/' \
	-e 's/must_have_one_flag+=("\(--.*\)=")/must_have_one_flag+=("\1")/' \
	-e "s/${LWORD}_filedir${RWORD}/__kubeadm_filedir/g" \
	-e "s/${LWORD}_get_comp_words_by_ref${RWORD}/__kubeadm_get_comp_words_by_ref/g" \
	-e "s/${LWORD}__ltrim_colon_completions${RWORD}/__kubeadm_ltrim_colon_completions/g" \
	-e "s/${LWORD}compgen${RWORD}/__kubeadm_compgen/g" \
	-e "s/${LWORD}compopt${RWORD}/__kubeadm_compopt/g" \
	-e "s/${LWORD}declare${RWORD}/builtin declare/g" \
	-e "s/\\\$(type${RWORD}/\$(__kubeadm_type/g" \
	<<'BASH_COMPLETION_EOF'
`
	out.Write([]byte(zsh_initialization))

	buf := new(bytes.Buffer)
	kubeadm.GenBashCompletion(buf)
	out.Write(buf.Bytes())

	zsh_tail := `
BASH_COMPLETION_EOF
}

__kubeadm_bash_source <(__kubeadm_convert_bash_to_zsh)
`
	out.Write([]byte(zsh_tail))
	return nil
}
<|endoftext|>


<commit_before>package orm

import (
	"strconv"
	"time"

	"github.com/go-pg/pg/types"
)

type URLValues map[string][]string

func (v URLValues) Has(name string) bool {
	_, ok := v[name]
	return ok
}

func (v URLValues) SetDefault(name string, values ...string) {
	if !v.Has(name) {
		v[name] = values
	}
}

func (v URLValues) Strings(name string) []string {
	return v[name]
}

func (v URLValues) String(name string) string {
	values := v.Strings(name)
	if len(values) == 0 {
		return ""
	}
	return values[0]
}

func (v URLValues) Int(name string) (int, error) {
	return strconv.Atoi(v.String(name))
}

func (v URLValues) Int64(name string) (int64, error) {
	return strconv.ParseInt(v.String(name), 10, 64)
}

func (v URLValues) Time(name string) (time.Time, error) {
	n, err := v.Int64(name)
	if err == nil {
		return time.Unix(n, 0), nil
	}
	return types.ParseTimeString(v.String(name))
}

func (v URLValues) Duration(name string) (time.Duration, error) {
	return time.ParseDuration(v.String(name))
}
<commit_msg>orm: add Maybe* helpers<commit_after>package orm

import (
	"net/url"
	"strconv"
	"time"

	"github.com/go-pg/pg/types"
)

type URLValues map[string][]string

func (v URLValues) Has(name string) bool {
	_, ok := v[name]
	return ok
}

func (v URLValues) SetDefault(name string, values ...string) {
	if !v.Has(name) {
		v[name] = values
	}
}

func (v URLValues) Strings(name string) []string {
	return v[name]
}

func (v URLValues) String(name string) string {
	values := v.Strings(name)
	if len(values) == 0 {
		return ""
	}
	return values[0]
}

func (v URLValues) Int(name string) (int, error) {
	return strconv.Atoi(v.String(name))
}

func (v URLValues) MaybeInt(name string) int {
	n, _ := v.Int(name)
	return n
}

func (v URLValues) Int64(name string) (int64, error) {
	return strconv.ParseInt(v.String(name), 10, 64)
}

func (v URLValues) MaybeInt64(name string) int64 {
	n, _ := v.Int64(name)
	return n
}

func (v URLValues) Time(name string) (time.Time, error) {
	n, err := v.Int64(name)
	if err == nil {
		return time.Unix(n, 0), nil
	}
	return types.ParseTimeString(v.String(name))
}

func (v URLValues) MaybeTime(name string) time.Time {
	tm, _ := v.Time(name)
	return tm
}

func (v URLValues) Duration(name string) (time.Duration, error) {
	return time.ParseDuration(v.String(name))
}

func (v URLValues) MaybeDuration(name string) time.Duration {
	dur, _ := v.Duration(name)
	return dur
}

func (v URLValues) Pager() *Pager {
	return NewPager(url.Values(v))
}
<|endoftext|>


<commit_before>package warehouse

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/lib/pq"
	"github.com/nishanths/fullstory"

	"github.com/fullstorydev/hauser/config"
)

type Redshift struct {
	conn         *sql.DB
	conf         *config.Config
	exportSchema Schema
	syncSchema   Schema
}

var (
	RedshiftTypeMap = FieldTypeMapper{
		"int64":     "BIGINT",
		"string":    "VARCHAR(max)",
		"time.Time": "TIMESTAMP",
	}
	beginningOfTime = time.Date(2015, 01, 01, 0, 0, 0, 0, time.UTC)
)

func NewRedshift(c *config.Config) *Redshift {
	if c.S3.S3Only {
		log.Printf("Config flag S3Only is on, data will not be loaded to Redshift")
	}
	return &Redshift{
		conf:         c,
		exportSchema: ExportTableSchema(RedshiftTypeMap),
		syncSchema:   SyncTableSchema(RedshiftTypeMap),
	}
}

// GetExportTableColumns returns all the columns of the export table.
// It opens a connection and calls getTableColumns
func (rs *Redshift) GetExportTableColumns() []string {
	var err error
	rs.conn, err = rs.MakeRedshfitConnection()
	if err != nil {
		log.Fatal(err)
	}
	defer rs.conn.Close()

	return rs.getTableColumns(rs.conf.Redshift.ExportTable)
}

func (rs *Redshift) ValueToString(val interface{}, isTime bool) string {
	s := fmt.Sprintf("%v", val)
	if isTime {
		t, _ := time.Parse(time.RFC3339Nano, s)
		return t.String()
	}

	s = strings.Replace(s, "\n", " ", -1)
	s = strings.Replace(s, "\r", " ", -1)
	s = strings.Replace(s, "\x00", "", -1)

	if len(s) >= rs.conf.Redshift.VarCharMax {
		s = s[:rs.conf.Redshift.VarCharMax-1]
	}
	return s
}

func (rs *Redshift) MakeRedshfitConnection() (*sql.DB, error) {
	url := fmt.Sprintf("user=%v password=%v host=%v port=%v dbname=%v",
		rs.conf.Redshift.User,
		rs.conf.Redshift.Password,
		rs.conf.Redshift.Host,
		rs.conf.Redshift.Port,
		rs.conf.Redshift.DB)

	var err error
	var db *sql.DB
	if db, err = sql.Open("postgres", url); err != nil {
		return nil, fmt.Errorf("redshift connect error : (%v)", err)
	}

	if err = db.Ping(); err != nil {
		return nil, fmt.Errorf("redshift ping error : (%v)", err)
	}
	return db, nil
}

func (rs *Redshift) UploadFile(name string) (string, error) {
	file, err := os.Open(name)
	if err != nil {
		return "", err
	}

	sess := session.Must(session.NewSession())
	svc := s3.New(sess, aws.NewConfig().WithRegion(rs.conf.S3.Region))

	ctx := context.Background()
	ctx, cancelFn := context.WithTimeout(ctx, rs.conf.S3.Timeout.Duration)
	defer cancelFn()

	_, objName := filepath.Split(name)
	s3path := fmt.Sprintf("s3://%s/%s", rs.conf.S3.Bucket, objName)

	_, err = svc.PutObjectWithContext(ctx, &s3.PutObjectInput{
		Bucket: aws.String(rs.conf.S3.Bucket),
		Key:    aws.String(objName),
		Body:   file,
	})

	return s3path, err
}

func (rs *Redshift) DeleteFile(s3obj string) {
	sess := session.Must(session.NewSession())
	svc := s3.New(sess, aws.NewConfig().WithRegion(rs.conf.S3.Region))

	ctx := context.Background()
	ctx, cancelFn := context.WithTimeout(ctx, rs.conf.S3.Timeout.Duration)
	defer cancelFn()

	_, objName := filepath.Split(s3obj)
	_, err := svc.DeleteObjectWithContext(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(rs.conf.S3.Bucket),
		Key:    aws.String(objName),
	})

	if err != nil {
		log.Printf("failed to delete S3 object %s: %s", s3obj, err)
		// just return - object will remain in S3
	}
}

func (rs *Redshift) LoadToWarehouse(s3obj string, _ ...fullstory.ExportMeta) error {
	var err error
	rs.conn, err = rs.MakeRedshfitConnection()
	if err != nil {
		return err
	}
	defer rs.conn.Close()

	if err = rs.CopyInData(s3obj); err != nil {
		return err
	}

	return nil
}

// EnsureCompatibleExportTable makes sure the export table has all the hauser schema columns
func (rs *Redshift) EnsureCompatibleExportTable() error {
	var err error
	rs.conn, err = rs.MakeRedshfitConnection()
	if err != nil {
		return err
	}
	defer rs.conn.Close()

	if !rs.DoesTableExist(rs.conf.Redshift.ExportTable) {
		// if the export table does not exist we create one with all the columns we expect!
		log.Printf("Export table %s does not exist! Creating one!", rs.conf.Redshift.ExportTable)
		if err = rs.CreateExportTable(); err != nil {
			return err
		}
		return nil
	}

	// make sure all the columns in the csv export exist in the Export table
	exportTableColumns := rs.getTableColumns(rs.conf.Redshift.ExportTable)
	missingFields := rs.getMissingFields(rs.exportSchema, exportTableColumns)

	// If some fields are missing from the fsexport table, either we added new fields
	// or existing expected columns were deleted by the user we add the relevant columns.
	// Alter the table and add the missing columns.
	if len(missingFields) > 0 {
		log.Printf("Found %d missing fields. Adding columns for these fields.", len(missingFields))
		for _, f := range missingFields {
			// Redshift only allows addition of one column at a time, hence the the alter statements in a loop yuck
			alterStmt := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s;", rs.conf.Redshift.ExportTable, f.Name, f.DBType)
			if _, err = rs.conn.Exec(alterStmt); err != nil {
				return err
			}
		}
	}
	return nil
}

// CopyInData copies data from the given s3File to the export table
func (rs *Redshift) CopyInData(s3file string) error {
	copy := fmt.Sprintf("COPY %s FROM '%s' CREDENTIALS '%s' DELIMITER ',' REGION '%s' FORMAT AS CSV ACCEPTINVCHARS;",
		rs.conf.Redshift.ExportTable, s3file, rs.conf.Redshift.Credentials, rs.conf.S3.Region)
	_, err := rs.conn.Exec(copy)
	return err
}

// CreateExportTable creates an export table with the hauser export table schema
func (rs *Redshift) CreateExportTable() error {
	log.Printf("Creating table %s", rs.conf.Redshift.ExportTable)

	stmt := fmt.Sprintf("create table %s(%s);", rs.conf.Redshift.ExportTable, rs.exportSchema.String())
	_, err := rs.conn.Exec(stmt)
	return err
}

// CreateSyncTable creates a sync table with the hauser sync table schema
func (rs *Redshift) CreateSyncTable() error {
	log.Printf("Creating table %s", rs.conf.Redshift.SyncTable)

	stmt := fmt.Sprintf("create table %s(%s);", rs.conf.Redshift.SyncTable, rs.syncSchema.String())
	_, err := rs.conn.Exec(stmt)
	return err
}

func (rs *Redshift) SaveSyncPoints(bundles ...fullstory.ExportMeta) error {
	var err error
	rs.conn, err = rs.MakeRedshfitConnection()
	if err != nil {
		log.Printf("Couldn't connect to DB: %s", err)
		return err
	}
	defer rs.conn.Close()

	for _, e := range bundles {
		insert := fmt.Sprintf("insert into %s values (%d, '%s', '%s')",
			rs.conf.Redshift.SyncTable, e.ID, time.Now().Format(time.RFC3339), e.Stop.Format(time.RFC3339))
		if _, err := rs.conn.Exec(insert); err != nil {
			return err
		}
	}

	return nil
}

func (rs *Redshift) DeleteExportRecordsAfter(end time.Time) error {
	stmt := fmt.Sprintf("DELETE FROM %s where EventStart > '%s';",
		rs.conf.Redshift.ExportTable, end.Format(time.RFC3339))
	_, err := rs.conn.Exec(stmt)
	if err != nil {
		log.Printf("failed to delete from %s: %s", rs.conf.Redshift.ExportTable, err)
		return err
	}

	return nil
}

func (rs *Redshift) LastSyncPoint() (time.Time, error) {
	t := beginningOfTime
	var err error
	rs.conn, err = rs.MakeRedshfitConnection()
	if err != nil {
		log.Printf("Couldn't connect to DB: %s", err)
		return t, err
	}
	defer rs.conn.Close()

	if rs.DoesTableExist(rs.conf.Redshift.SyncTable) {
		var syncTime pq.NullTime
		q := fmt.Sprintf("SELECT max(BundleEndTime) FROM %s;", rs.conf.Redshift.SyncTable)
		if err := rs.conn.QueryRow(q).Scan(&syncTime); err != nil {
			log.Printf("Couldn't get max(BundleEndTime): %s", err)
			return t, err
		}
		if syncTime.Valid {
			t = syncTime.Time
		}

		if err := rs.RemoveOrphanedRecords(syncTime); err != nil {
			return t, err
		}

	} else {
		if err := rs.CreateSyncTable(); err != nil {
			log.Printf("Couldn't create sync table: %s", err)
			return t, err
		}
	}
	return t, nil
}

func (rs *Redshift) RemoveOrphanedRecords(lastSync pq.NullTime) error {
	if rs.conf.S3.S3Only {
		// no need to check for orphaned records, as we're not loading to the export table
		return nil
	}

	if !rs.DoesTableExist(rs.conf.Redshift.ExportTable) {
		if err := rs.CreateExportTable(); err != nil {
			log.Printf("Couldn't create export table: %s", err)
			return err
		}
	}

	// Find the time of the latest export record...if it's after
	// the time in the sync table, then there must have been a failure
	// after some records have been loaded, but before the sync record
	// was written. Use this as the latest sync time, and don't load
	// any records before this point to prevent duplication
	var exportTime pq.NullTime
	q := fmt.Sprintf("SELECT max(EventStart) FROM %s;", rs.conf.Redshift.ExportTable)
	if err := rs.conn.QueryRow(q).Scan(&exportTime); err != nil {
		log.Printf("Couldn't get max(EventStart): %s", err)
		return err
	}
	if exportTime.Valid && exportTime.Time.After(lastSync.Time) {
		log.Printf("Export record timestamp after sync time (%s vs %s); cleaning",
			exportTime.Time, lastSync.Time)
		rs.DeleteExportRecordsAfter(lastSync.Time)
	}

	return nil
}

// DoesTableExist checks if a table with a given name exists in the public schema
func (rs *Redshift) DoesTableExist(name string) bool {
	log.Printf("Checking if table %s exists", name)

	var exists int
	err := rs.conn.QueryRow("SELECT count(*) FROM pg_tables WHERE schemaname = 'public' AND tablename = $1;", name).Scan(&exists)
	if err != nil {
		// something is horribly wrong...just give up
		log.Fatal(err)
	}
	return (exists != 0)
}

func (rs *Redshift) getTableColumns(name string) []string {
	ctx := context.Background()
	log.Printf("Fetching columns for table %s", name)
	rows, err := rs.conn.QueryContext(ctx,
		"SELECT \"column\" FROM pg_table_def WHERE schemaname = 'public' AND tablename = $1;", name)
	if err != nil {
		log.Fatal(err)
	}
	var columns []string

	defer rows.Close()
	for rows.Next(){
		var column string
		err = rows.Scan(&column)
		columns = append(columns, column)
	}

	// get any error encountered during iteration
	if err = rows.Err(); err != nil {
		log.Fatal(err)
	}
	return columns
}

func (rs *Redshift) getMissingFields(schema Schema, tableColumns []string) []WarehouseField {
	existingColumns := make(map[string]struct{})
	for _, column := range tableColumns {
		// Redshift columns are case insensitive
		existingColumns[strings.ToLower(column)] = struct{}{}
	}

	var missingFields []WarehouseField
	for _, f := range schema {
		if _, ok := existingColumns[strings.ToLower(f.Name)]; !ok {
			missingFields = append(missingFields, f)
		}
	}

	return missingFields
}

func (rs *Redshift) GetUploadFailedMsg(filename string, err error) string {
	return fmt.Sprintf("Failed to upload file %s to s3: %s", filename, err)
}

func (rs *Redshift) IsUploadOnly() bool {
	return rs.conf.S3.S3Only
}
<commit_msg>MakeRedshiftConnection spelling<commit_after>package warehouse

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/lib/pq"
	"github.com/nishanths/fullstory"

	"github.com/fullstorydev/hauser/config"
)

type Redshift struct {
	conn         *sql.DB
	conf         *config.Config
	exportSchema Schema
	syncSchema   Schema
}

var (
	RedshiftTypeMap = FieldTypeMapper{
		"int64":     "BIGINT",
		"string":    "VARCHAR(max)",
		"time.Time": "TIMESTAMP",
	}
	beginningOfTime = time.Date(2015, 01, 01, 0, 0, 0, 0, time.UTC)
)

func NewRedshift(c *config.Config) *Redshift {
	if c.S3.S3Only {
		log.Printf("Config flag S3Only is on, data will not be loaded to Redshift")
	}
	return &Redshift{
		conf:         c,
		exportSchema: ExportTableSchema(RedshiftTypeMap),
		syncSchema:   SyncTableSchema(RedshiftTypeMap),
	}
}

// GetExportTableColumns returns all the columns of the export table.
// It opens a connection and calls getTableColumns
func (rs *Redshift) GetExportTableColumns() []string {
	var err error
	rs.conn, err = rs.MakeRedshiftConnection()
	if err != nil {
		log.Fatal(err)
	}
	defer rs.conn.Close()

	return rs.getTableColumns(rs.conf.Redshift.ExportTable)
}

func (rs *Redshift) ValueToString(val interface{}, isTime bool) string {
	s := fmt.Sprintf("%v", val)
	if isTime {
		t, _ := time.Parse(time.RFC3339Nano, s)
		return t.String()
	}

	s = strings.Replace(s, "\n", " ", -1)
	s = strings.Replace(s, "\r", " ", -1)
	s = strings.Replace(s, "\x00", "", -1)

	if len(s) >= rs.conf.Redshift.VarCharMax {
		s = s[:rs.conf.Redshift.VarCharMax-1]
	}
	return s
}

func (rs *Redshift) MakeRedshiftConnection() (*sql.DB, error) {
	url := fmt.Sprintf("user=%v password=%v host=%v port=%v dbname=%v",
		rs.conf.Redshift.User,
		rs.conf.Redshift.Password,
		rs.conf.Redshift.Host,
		rs.conf.Redshift.Port,
		rs.conf.Redshift.DB)

	var err error
	var db *sql.DB
	if db, err = sql.Open("postgres", url); err != nil {
		return nil, fmt.Errorf("redshift connect error : (%v)", err)
	}

	if err = db.Ping(); err != nil {
		return nil, fmt.Errorf("redshift ping error : (%v)", err)
	}
	return db, nil
}

func (rs *Redshift) UploadFile(name string) (string, error) {
	file, err := os.Open(name)
	if err != nil {
		return "", err
	}

	sess := session.Must(session.NewSession())
	svc := s3.New(sess, aws.NewConfig().WithRegion(rs.conf.S3.Region))

	ctx := context.Background()
	ctx, cancelFn := context.WithTimeout(ctx, rs.conf.S3.Timeout.Duration)
	defer cancelFn()

	_, objName := filepath.Split(name)
	s3path := fmt.Sprintf("s3://%s/%s", rs.conf.S3.Bucket, objName)

	_, err = svc.PutObjectWithContext(ctx, &s3.PutObjectInput{
		Bucket: aws.String(rs.conf.S3.Bucket),
		Key:    aws.String(objName),
		Body:   file,
	})

	return s3path, err
}

func (rs *Redshift) DeleteFile(s3obj string) {
	sess := session.Must(session.NewSession())
	svc := s3.New(sess, aws.NewConfig().WithRegion(rs.conf.S3.Region))

	ctx := context.Background()
	ctx, cancelFn := context.WithTimeout(ctx, rs.conf.S3.Timeout.Duration)
	defer cancelFn()

	_, objName := filepath.Split(s3obj)
	_, err := svc.DeleteObjectWithContext(ctx, &s3.DeleteObjectInput{
		Bucket: aws.String(rs.conf.S3.Bucket),
		Key:    aws.String(objName),
	})

	if err != nil {
		log.Printf("failed to delete S3 object %s: %s", s3obj, err)
		// just return - object will remain in S3
	}
}

func (rs *Redshift) LoadToWarehouse(s3obj string, _ ...fullstory.ExportMeta) error {
	var err error
	rs.conn, err = rs.MakeRedshiftConnection()
	if err != nil {
		return err
	}
	defer rs.conn.Close()

	if err = rs.CopyInData(s3obj); err != nil {
		return err
	}

	return nil
}

// EnsureCompatibleExportTable makes sure the export table has all the hauser schema columns
func (rs *Redshift) EnsureCompatibleExportTable() error {
	var err error
	rs.conn, err = rs.MakeRedshiftConnection()
	if err != nil {
		return err
	}
	defer rs.conn.Close()

	if !rs.DoesTableExist(rs.conf.Redshift.ExportTable) {
		// if the export table does not exist we create one with all the columns we expect!
		log.Printf("Export table %s does not exist! Creating one!", rs.conf.Redshift.ExportTable)
		if err = rs.CreateExportTable(); err != nil {
			return err
		}
		return nil
	}

	// make sure all the columns in the csv export exist in the Export table
	exportTableColumns := rs.getTableColumns(rs.conf.Redshift.ExportTable)
	missingFields := rs.getMissingFields(rs.exportSchema, exportTableColumns)

	// If some fields are missing from the fsexport table, either we added new fields
	// or existing expected columns were deleted by the user we add the relevant columns.
	// Alter the table and add the missing columns.
	if len(missingFields) > 0 {
		log.Printf("Found %d missing fields. Adding columns for these fields.", len(missingFields))
		for _, f := range missingFields {
			// Redshift only allows addition of one column at a time, hence the the alter statements in a loop yuck
			alterStmt := fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s;", rs.conf.Redshift.ExportTable, f.Name, f.DBType)
			if _, err = rs.conn.Exec(alterStmt); err != nil {
				return err
			}
		}
	}
	return nil
}

// CopyInData copies data from the given s3File to the export table
func (rs *Redshift) CopyInData(s3file string) error {
	copy := fmt.Sprintf("COPY %s FROM '%s' CREDENTIALS '%s' DELIMITER ',' REGION '%s' FORMAT AS CSV ACCEPTINVCHARS;",
		rs.conf.Redshift.ExportTable, s3file, rs.conf.Redshift.Credentials, rs.conf.S3.Region)
	_, err := rs.conn.Exec(copy)
	return err
}

// CreateExportTable creates an export table with the hauser export table schema
func (rs *Redshift) CreateExportTable() error {
	log.Printf("Creating table %s", rs.conf.Redshift.ExportTable)

	stmt := fmt.Sprintf("create table %s(%s);", rs.conf.Redshift.ExportTable, rs.exportSchema.String())
	_, err := rs.conn.Exec(stmt)
	return err
}

// CreateSyncTable creates a sync table with the hauser sync table schema
func (rs *Redshift) CreateSyncTable() error {
	log.Printf("Creating table %s", rs.conf.Redshift.SyncTable)

	stmt := fmt.Sprintf("create table %s(%s);", rs.conf.Redshift.SyncTable, rs.syncSchema.String())
	_, err := rs.conn.Exec(stmt)
	return err
}

func (rs *Redshift) SaveSyncPoints(bundles ...fullstory.ExportMeta) error {
	var err error
	rs.conn, err = rs.MakeRedshiftConnection()
	if err != nil {
		log.Printf("Couldn't connect to DB: %s", err)
		return err
	}
	defer rs.conn.Close()

	for _, e := range bundles {
		insert := fmt.Sprintf("insert into %s values (%d, '%s', '%s')",
			rs.conf.Redshift.SyncTable, e.ID, time.Now().Format(time.RFC3339), e.Stop.Format(time.RFC3339))
		if _, err := rs.conn.Exec(insert); err != nil {
			return err
		}
	}

	return nil
}

func (rs *Redshift) DeleteExportRecordsAfter(end time.Time) error {
	stmt := fmt.Sprintf("DELETE FROM %s where EventStart > '%s';",
		rs.conf.Redshift.ExportTable, end.Format(time.RFC3339))
	_, err := rs.conn.Exec(stmt)
	if err != nil {
		log.Printf("failed to delete from %s: %s", rs.conf.Redshift.ExportTable, err)
		return err
	}

	return nil
}

func (rs *Redshift) LastSyncPoint() (time.Time, error) {
	t := beginningOfTime
	var err error
	rs.conn, err = rs.MakeRedshiftConnection()
	if err != nil {
		log.Printf("Couldn't connect to DB: %s", err)
		return t, err
	}
	defer rs.conn.Close()

	if rs.DoesTableExist(rs.conf.Redshift.SyncTable) {
		var syncTime pq.NullTime
		q := fmt.Sprintf("SELECT max(BundleEndTime) FROM %s;", rs.conf.Redshift.SyncTable)
		if err := rs.conn.QueryRow(q).Scan(&syncTime); err != nil {
			log.Printf("Couldn't get max(BundleEndTime): %s", err)
			return t, err
		}
		if syncTime.Valid {
			t = syncTime.Time
		}

		if err := rs.RemoveOrphanedRecords(syncTime); err != nil {
			return t, err
		}

	} else {
		if err := rs.CreateSyncTable(); err != nil {
			log.Printf("Couldn't create sync table: %s", err)
			return t, err
		}
	}
	return t, nil
}

func (rs *Redshift) RemoveOrphanedRecords(lastSync pq.NullTime) error {
	if rs.conf.S3.S3Only {
		// no need to check for orphaned records, as we're not loading to the export table
		return nil
	}

	if !rs.DoesTableExist(rs.conf.Redshift.ExportTable) {
		if err := rs.CreateExportTable(); err != nil {
			log.Printf("Couldn't create export table: %s", err)
			return err
		}
	}

	// Find the time of the latest export record...if it's after
	// the time in the sync table, then there must have been a failure
	// after some records have been loaded, but before the sync record
	// was written. Use this as the latest sync time, and don't load
	// any records before this point to prevent duplication
	var exportTime pq.NullTime
	q := fmt.Sprintf("SELECT max(EventStart) FROM %s;", rs.conf.Redshift.ExportTable)
	if err := rs.conn.QueryRow(q).Scan(&exportTime); err != nil {
		log.Printf("Couldn't get max(EventStart): %s", err)
		return err
	}
	if exportTime.Valid && exportTime.Time.After(lastSync.Time) {
		log.Printf("Export record timestamp after sync time (%s vs %s); cleaning",
			exportTime.Time, lastSync.Time)
		rs.DeleteExportRecordsAfter(lastSync.Time)
	}

	return nil
}

// DoesTableExist checks if a table with a given name exists in the public schema
func (rs *Redshift) DoesTableExist(name string) bool {
	log.Printf("Checking if table %s exists", name)

	var exists int
	err := rs.conn.QueryRow("SELECT count(*) FROM pg_tables WHERE schemaname = 'public' AND tablename = $1;", name).Scan(&exists)
	if err != nil {
		// something is horribly wrong...just give up
		log.Fatal(err)
	}
	return (exists != 0)
}

func (rs *Redshift) getTableColumns(name string) []string {
	ctx := context.Background()
	log.Printf("Fetching columns for table %s", name)
	rows, err := rs.conn.QueryContext(ctx,
		"SELECT \"column\" FROM pg_table_def WHERE schemaname = 'public' AND tablename = $1;", name)
	if err != nil {
		log.Fatal(err)
	}
	var columns []string

	defer rows.Close()
	for rows.Next(){
		var column string
		err = rows.Scan(&column)
		columns = append(columns, column)
	}

	// get any error encountered during iteration
	if err = rows.Err(); err != nil {
		log.Fatal(err)
	}
	return columns
}

func (rs *Redshift) getMissingFields(schema Schema, tableColumns []string) []WarehouseField {
	existingColumns := make(map[string]struct{})
	for _, column := range tableColumns {
		// Redshift columns are case insensitive
		existingColumns[strings.ToLower(column)] = struct{}{}
	}

	var missingFields []WarehouseField
	for _, f := range schema {
		if _, ok := existingColumns[strings.ToLower(f.Name)]; !ok {
			missingFields = append(missingFields, f)
		}
	}

	return missingFields
}

func (rs *Redshift) GetUploadFailedMsg(filename string, err error) string {
	return fmt.Sprintf("Failed to upload file %s to s3: %s", filename, err)
}

func (rs *Redshift) IsUploadOnly() bool {
	return rs.conf.S3.S3Only
}
<|endoftext|>


<commit_before>package web

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/eyecuelab/kit/flect"
	"github.com/google/jsonapi"
	"github.com/jinzhu/inflection"
	"github.com/spf13/viper"
)

// JSONApiAction jsonapi action
type JSONApiAction struct {
	Method string         `json:"method"`
	Name   string         `json:"name"`
	URL    string         `json:"url"`
	Fields []JSONApiField `json:"fields"`
}

// JSONApiField jsonapi action field
type JSONApiField struct {
	Name      string `json:"name"`
	InputType string `json:"type"`
	Value     string `json:"value"`
	Required  bool   `json:"required"`
}

// APIURL full api url for the path
func APIURL(path string) string {
	return fmt.Sprintf("%s%s", viper.GetString("root_url"), path)
}

// JSONApiSelfLink self link helper
func JSONApiSelfLink(i interface{}) *jsonapi.Links {
	name := reflect.TypeOf(i).Name()
	name = inflection.Plural(strings.ToLower(name))

	id := reflect.ValueOf(i).FieldByName("ID").Interface()

	return &jsonapi.Links{
		"self": jsonapi.Link{
			Href: fmt.Sprintf("/%v/%v", name, id)},
	}
}

// JSONApiRefLink ref link helper
func JSONApiRefLink(i interface{}, rel string) *jsonapi.Links {
	fields := flect.Fields(reflect.ValueOf(i))
	name := reflect.TypeOf(i).Name()
	id := reflect.ValueOf(i).FieldByName("ID").Interface()
	for _, f := range fields {
		value, opts, ok := flect.TagValues(f.Tag, "jsonapi")
		if ok && value == "relation" && opts.HasOption(rel) {
			if _, _, ok := flect.TagValues(f.Tag, "link"); ok {
				return &jsonapi.Links{
					"related": fmt.Sprintf(
						"/%v/%v/%v",
						strings.ToLower(inflection.Plural(name)),
						id,
						strings.ToLower(inflection.Plural(f.Name)),
					),
				}
			}
		}
	}
	return &jsonapi.Links{}
}
<commit_msg>Update jsonapi action and field types<commit_after>package web

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/eyecuelab/kit/flect"
	"github.com/google/jsonapi"
	"github.com/jinzhu/inflection"
	"github.com/spf13/viper"
)

// JSONApiAction jsonapi action
type JSONApiAction struct {
	Method string         `json:"method"`
	Name   string         `json:"name"`
	ID     string         `json:"id"`
	Type   string         `json:"type"`
	URL    string         `json:"url"`
	Fields []JSONApiField `json:"fields"`
}

// JSONApiField jsonapi action field
type JSONApiField struct {
	Name      string      `json:"name"`
	InputType string      `json:"type"`
	Value     interface{} `json:"value"`
	Required  bool        `json:"required"`
}

// APIURL full api url for the path
func APIURL(path string) string {
	return fmt.Sprintf("%s%s", viper.GetString("root_url"), path)
}

// JSONApiSelfLink self link helper
func JSONApiSelfLink(i interface{}) *jsonapi.Links {
	name := reflect.TypeOf(i).Name()
	name = inflection.Plural(strings.ToLower(name))

	id := reflect.ValueOf(i).FieldByName("ID").Interface()

	return &jsonapi.Links{
		"self": jsonapi.Link{
			Href: fmt.Sprintf("/%v/%v", name, id)},
	}
}

// JSONApiRefLink ref link helper
func JSONApiRefLink(i interface{}, rel string) *jsonapi.Links {
	fields := flect.Fields(reflect.ValueOf(i))
	name := reflect.TypeOf(i).Name()
	id := reflect.ValueOf(i).FieldByName("ID").Interface()
	for _, f := range fields {
		value, opts, ok := flect.TagValues(f.Tag, "jsonapi")
		if ok && value == "relation" && opts.HasOption(rel) {
			if _, _, ok := flect.TagValues(f.Tag, "link"); ok {
				return &jsonapi.Links{
					"related": fmt.Sprintf(
						"/%v/%v/%v",
						strings.ToLower(inflection.Plural(name)),
						id,
						strings.ToLower(inflection.Plural(f.Name)),
					),
				}
			}
		}
	}
	return &jsonapi.Links{}
}
<|endoftext|>


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

package confmaptest // import "go.opentelemetry.io/collector/confmap/confmaptest"

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"

	"gopkg.in/yaml.v2"

	"go.opentelemetry.io/collector/confmap"
)

// LoadConf loads a confmap.Conf from file, and does NOT validate the configuration.
func LoadConf(fileName string) (*confmap.Conf, error) {
	// Clean the path before using it.
	content, err := os.ReadFile(filepath.Clean(fileName))
	if err != nil {
		return nil, fmt.Errorf("unable to read the file %v: %w", fileName, err)
	}

	var rawConf map[string]interface{}
	if err = yaml.Unmarshal(content, &rawConf); err != nil {
		return nil, err
	}

	return confmap.NewFromStringMap(rawConf), nil
}

var schemeValidator = regexp.MustCompile("^[A-Za-z][A-Za-z0-9+.-]+$")

// ValidateProviderScheme enforces that given confmap.Provider.Scheme() object is following the restriction defined by the collector:
//   - Checks that the scheme name follows the restrictions defined https://datatracker.ietf.org/doc/html/rfc3986#section-3.1
//   - Checks that the scheme name has at leas two characters per the confmap.Provider.Scheme() comment.
func ValidateProviderScheme(p confmap.Provider) error {
	scheme := p.Scheme()
	if len(scheme) < 2 {
		return fmt.Errorf("scheme must be at least 2 characters long: %q", scheme)
	}

	if !schemeValidator.MatchString(scheme) {
		return fmt.Errorf("scheme names consist of a sequence of characters beginning with a letter and followed by any combination of letters, digits, \"+\", \".\", or \"-\": %q", scheme)
	}

	return nil
}
<commit_msg>[chore] fix wrong usage of yaml.v2 instead of yaml.v3 (#5892)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package confmaptest // import "go.opentelemetry.io/collector/confmap/confmaptest"

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"

	"gopkg.in/yaml.v3"

	"go.opentelemetry.io/collector/confmap"
)

// LoadConf loads a confmap.Conf from file, and does NOT validate the configuration.
func LoadConf(fileName string) (*confmap.Conf, error) {
	// Clean the path before using it.
	content, err := os.ReadFile(filepath.Clean(fileName))
	if err != nil {
		return nil, fmt.Errorf("unable to read the file %v: %w", fileName, err)
	}

	var rawConf map[string]interface{}
	if err = yaml.Unmarshal(content, &rawConf); err != nil {
		return nil, err
	}

	return confmap.NewFromStringMap(rawConf), nil
}

var schemeValidator = regexp.MustCompile("^[A-Za-z][A-Za-z0-9+.-]+$")

// ValidateProviderScheme enforces that given confmap.Provider.Scheme() object is following the restriction defined by the collector:
//   - Checks that the scheme name follows the restrictions defined https://datatracker.ietf.org/doc/html/rfc3986#section-3.1
//   - Checks that the scheme name has at leas two characters per the confmap.Provider.Scheme() comment.
func ValidateProviderScheme(p confmap.Provider) error {
	scheme := p.Scheme()
	if len(scheme) < 2 {
		return fmt.Errorf("scheme must be at least 2 characters long: %q", scheme)
	}

	if !schemeValidator.MatchString(scheme) {
		return fmt.Errorf("scheme names consist of a sequence of characters beginning with a letter and followed by any combination of letters, digits, \"+\", \".\", or \"-\": %q", scheme)
	}

	return nil
}
<|endoftext|>


<commit_before>// +build !race

package command

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

	"github.com/hashicorp/vault/vault/diagnose"
	"github.com/mitchellh/cli"
)

func testOperatorDiagnoseCommand(tb testing.TB) *OperatorDiagnoseCommand {
	tb.Helper()

	ui := cli.NewMockUi()
	return &OperatorDiagnoseCommand{
		diagnose: diagnose.New(ioutil.Discard),
		BaseCommand: &BaseCommand{
			UI: ui,
		},
		skipEndEnd: true,
	}
}

func TestOperatorDiagnoseCommand_Run(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name     string
		args     []string
		expected []*diagnose.Result
	}{
		{
			"diagnose_ok",
			[]string{
				"-config", "./server/test-fixtures/config_diagnose_ok.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "operating system",
					Status: diagnose.OkStatus,
					Children: []*diagnose.Result{
						{
							Name:   "open file limits",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "disk usage",
							Status: diagnose.OkStatus,
						},
					},
				},

				{
					Name:   "parse-config",
					Status: diagnose.OkStatus,
				},
				{
					Name:   "init-listeners",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-listeners",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "check-listener-tls",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"TLS is disabled in a Listener config stanza.",
							},
						},
					},
				},
				{
					Name:   "storage",
					Status: diagnose.OkStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_storage",
			[]string{
				"-config", "./server/test-fixtures/nostore_config.hcl",
			},
			[]*diagnose.Result{
				{
					Name:    "storage",
					Status:  diagnose.ErrorStatus,
					Message: "no storage stanza found in config",
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.ErrorStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_listener_config_ok",
			[]string{
				"-config", "./server/test-fixtures/tls_config_ok.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "init-listeners",
					Status: diagnose.OkStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-listeners",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "check-listener-tls",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_https_storage",
			[]string{
				"-config", "./server/test-fixtures/config_bad_https_storage.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:    "test-storage-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "expired",
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_https_hastorage",
			[]string{
				"-config", "./server/test-fixtures/config_diagnose_hastorage_bad_https.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"consul storage does not connect to local agent, but directly to server",
							},
						},
					},
				},
				{
					Name:   "setup-ha-storage",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-ha-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"consul storage does not connect to local agent, but directly to server",
							},
						},
						{
							Name:    "test-ha-storage-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "x509: certificate has expired or is not yet valid",
						},
					},
				},
				{
					Name:   "find-cluster-addr",
					Status: diagnose.ErrorStatus,
				},
			},
		},
		{
			"diagnose_invalid_https_sr",
			[]string{
				"-config", "./server/test-fixtures/diagnose_bad_https_consul_sr.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "service-discovery",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:    "test-serviceregistration-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "failed to verify certificate: x509: certificate has expired or is not yet valid",
						},
						{
							Name:   "test-consul-direct-access-service-discovery",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								diagnose.DirAccessErr,
							},
						},
					},
				},
			},
		},
		{
			"diagnose_direct_storage_access",
			[]string{
				"-config", "./server/test-fixtures/diagnose_ok_storage_direct_access.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								diagnose.DirAccessErr,
							},
						},
					},
				},
			},
		},
	}

	t.Run("validations", func(t *testing.T) {
		t.Parallel()

		for _, tc := range cases {
			tc := tc
			t.Run(tc.name, func(t *testing.T) {
				t.Parallel()
				client, closer := testVaultServer(t)
				defer closer()

				cmd := testOperatorDiagnoseCommand(t)
				cmd.client = client

				cmd.Run(tc.args)
				result := cmd.diagnose.Finalize(context.Background())

				if err := compareResults(tc.expected, result.Children); err != nil {
					t.Fatalf("Did not find expected test results: %v", err)
					t.Fatal(result.String())
				}
			})
		}
	})
}

func compareResults(expected []*diagnose.Result, actual []*diagnose.Result) error {
	for _, exp := range expected {
		found := false
		// Check them all so we don't have to be order specific
		for _, act := range actual {
			if exp.Name == act.Name {
				found = true
				if err := compareResult(exp, act); err != nil {
					return err
				}
				break
			}
		}
		if !found {
			return fmt.Errorf("could not find expected test result: %s", exp.Name)
		}
	}
	return nil
}

func compareResult(exp *diagnose.Result, act *diagnose.Result) error {
	if exp.Name != act.Name {
		return fmt.Errorf("names mismatch: %s vs %s", exp.Name, act.Name)
	}
	if exp.Status != act.Status {
		if act.Status != diagnose.OkStatus {
			return fmt.Errorf("section %s, status mismatch: %s vs %s, got error %s", exp.Name, exp.Status, act.Status, act.Message)

		}
		return fmt.Errorf("section %s, status mismatch: %s vs %s", exp.Name, exp.Status, act.Status)
	}
	if exp.Message != "" && exp.Message != act.Message && !strings.Contains(act.Message, exp.Message) {
		return fmt.Errorf("section %s, message not found: %s in %s", exp.Name, exp.Message, act.Message)
	}
	if len(exp.Warnings) != len(act.Warnings) {
		return fmt.Errorf("section %s, warning count mismatch: %d vs %d", exp.Name, len(exp.Warnings), len(act.Warnings))
	}
	for j := range exp.Warnings {
		if !strings.Contains(act.Warnings[j], exp.Warnings[j]) {
			return fmt.Errorf("section %s, warning message not found: %s in %s", exp.Name, exp.Warnings[j], act.Warnings[j])
		}
	}
	if len(exp.Children) > len(act.Children) {
		errStrings := []string{}
		for _, c := range act.Children {
			errStrings = append(errStrings, fmt.Sprintf("%+v", c))
		}
		return fmt.Errorf(strings.Join(errStrings, ","))
	}

	if len(exp.Children) > 0 {
		return compareResults(exp.Children, act.Children)
	}

	return nil
}
<commit_msg>remove operating system stanza from tests on account of disk usage being different across ci for different runs (#11764)<commit_after>// +build !race

package command

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

	"github.com/hashicorp/vault/vault/diagnose"
	"github.com/mitchellh/cli"
)

func testOperatorDiagnoseCommand(tb testing.TB) *OperatorDiagnoseCommand {
	tb.Helper()

	ui := cli.NewMockUi()
	return &OperatorDiagnoseCommand{
		diagnose: diagnose.New(ioutil.Discard),
		BaseCommand: &BaseCommand{
			UI: ui,
		},
		skipEndEnd: true,
	}
}

func TestOperatorDiagnoseCommand_Run(t *testing.T) {
	t.Parallel()
	cases := []struct {
		name     string
		args     []string
		expected []*diagnose.Result
	}{
		{
			"diagnose_ok",
			[]string{
				"-config", "./server/test-fixtures/config_diagnose_ok.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "parse-config",
					Status: diagnose.OkStatus,
				},
				{
					Name:   "init-listeners",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-listeners",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "check-listener-tls",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"TLS is disabled in a Listener config stanza.",
							},
						},
					},
				},
				{
					Name:   "storage",
					Status: diagnose.OkStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_storage",
			[]string{
				"-config", "./server/test-fixtures/nostore_config.hcl",
			},
			[]*diagnose.Result{
				{
					Name:    "storage",
					Status:  diagnose.ErrorStatus,
					Message: "no storage stanza found in config",
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.ErrorStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_listener_config_ok",
			[]string{
				"-config", "./server/test-fixtures/tls_config_ok.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "init-listeners",
					Status: diagnose.OkStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-listeners",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "check-listener-tls",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_https_storage",
			[]string{
				"-config", "./server/test-fixtures/config_bad_https_storage.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:    "test-storage-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "expired",
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.OkStatus,
						},
					},
				},
			},
		},
		{
			"diagnose_invalid_https_hastorage",
			[]string{
				"-config", "./server/test-fixtures/config_diagnose_hastorage_bad_https.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"consul storage does not connect to local agent, but directly to server",
							},
						},
					},
				},
				{
					Name:   "setup-ha-storage",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-ha-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								"consul storage does not connect to local agent, but directly to server",
							},
						},
						{
							Name:    "test-ha-storage-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "x509: certificate has expired or is not yet valid",
						},
					},
				},
				{
					Name:   "find-cluster-addr",
					Status: diagnose.ErrorStatus,
				},
			},
		},
		{
			"diagnose_invalid_https_sr",
			[]string{
				"-config", "./server/test-fixtures/diagnose_bad_https_consul_sr.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "service-discovery",
					Status: diagnose.ErrorStatus,
					Children: []*diagnose.Result{
						{
							Name:    "test-serviceregistration-tls-consul",
							Status:  diagnose.ErrorStatus,
							Message: "failed to verify certificate: x509: certificate has expired or is not yet valid",
						},
						{
							Name:   "test-consul-direct-access-service-discovery",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								diagnose.DirAccessErr,
							},
						},
					},
				},
			},
		},
		{
			"diagnose_direct_storage_access",
			[]string{
				"-config", "./server/test-fixtures/diagnose_ok_storage_direct_access.hcl",
			},
			[]*diagnose.Result{
				{
					Name:   "storage",
					Status: diagnose.WarningStatus,
					Children: []*diagnose.Result{
						{
							Name:   "create-storage-backend",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-storage-tls-consul",
							Status: diagnose.OkStatus,
						},
						{
							Name:   "test-consul-direct-access-storage",
							Status: diagnose.WarningStatus,
							Warnings: []string{
								diagnose.DirAccessErr,
							},
						},
					},
				},
			},
		},
	}

	t.Run("validations", func(t *testing.T) {
		t.Parallel()

		for _, tc := range cases {
			tc := tc
			t.Run(tc.name, func(t *testing.T) {
				t.Parallel()
				client, closer := testVaultServer(t)
				defer closer()

				cmd := testOperatorDiagnoseCommand(t)
				cmd.client = client

				cmd.Run(tc.args)
				result := cmd.diagnose.Finalize(context.Background())

				if err := compareResults(tc.expected, result.Children); err != nil {
					t.Fatalf("Did not find expected test results: %v", err)
				}
			})
		}
	})
}

func compareResults(expected []*diagnose.Result, actual []*diagnose.Result) error {
	for _, exp := range expected {
		found := false
		// Check them all so we don't have to be order specific
		for _, act := range actual {
			if exp.Name == act.Name {
				found = true
				if err := compareResult(exp, act); err != nil {
					return err
				}
				break
			}
		}
		if !found {
			return fmt.Errorf("could not find expected test result: %s", exp.Name)
		}
	}
	return nil
}

func compareResult(exp *diagnose.Result, act *diagnose.Result) error {
	if exp.Name != act.Name {
		return fmt.Errorf("names mismatch: %s vs %s", exp.Name, act.Name)
	}
	if exp.Status != act.Status {
		if act.Status != diagnose.OkStatus {
			return fmt.Errorf("section %s, status mismatch: %s vs %s, got error %s", exp.Name, exp.Status, act.Status, act.Message)

		}
		return fmt.Errorf("section %s, status mismatch: %s vs %s", exp.Name, exp.Status, act.Status)
	}
	if exp.Message != "" && exp.Message != act.Message && !strings.Contains(act.Message, exp.Message) {
		return fmt.Errorf("section %s, message not found: %s in %s", exp.Name, exp.Message, act.Message)
	}
	if len(exp.Warnings) != len(act.Warnings) {
		return fmt.Errorf("section %s, warning count mismatch: %d vs %d", exp.Name, len(exp.Warnings), len(act.Warnings))
	}
	for j := range exp.Warnings {
		if !strings.Contains(act.Warnings[j], exp.Warnings[j]) {
			return fmt.Errorf("section %s, warning message not found: %s in %s", exp.Name, exp.Warnings[j], act.Warnings[j])
		}
	}
	if len(exp.Children) > len(act.Children) {
		errStrings := []string{}
		for _, c := range act.Children {
			errStrings = append(errStrings, fmt.Sprintf("%+v", c))
		}
		return fmt.Errorf(strings.Join(errStrings, ","))
	}

	if len(exp.Children) > 0 {
		return compareResults(exp.Children, act.Children)
	}

	return nil
}
<|endoftext|>


<commit_before>package common

import (
	"bytes"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"testing"
)

func TestStepCreateFloppy_Impl(t *testing.T) {
	var raw interface{}
	raw = new(StepCreateFloppy)
	if _, ok := raw.(multistep.Step); !ok {
		t.Fatalf("StepCreateFloppy should be a step")
	}
}

func testStepCreateFloppyState(t *testing.T) multistep.StateBag {
	state := new(multistep.BasicStateBag)
	state.Put("ui", &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	})
	return state
}

func TestStepCreateFloppy(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 10
	expected := count
	files := make([]string, count)

	prefix := "exists"
	ext := ".tmp"

	for i := 0; i < expected; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i)+ext)

		_, err := os.Create(files[i])
		if err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	lists := [][]string{
		files,
		[]string{dir + string(os.PathSeparator) + prefix + "*" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "?" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "[0123456789]" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "[0-9]" + ext},
		[]string{dir + string(os.PathSeparator)},
		[]string{dir},
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionContinue {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); ok {
			t.Fatalf("state should be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path").(string)

		if _, err := os.Stat(floppy_path); err != nil {
			t.Fatal("file not found: %s for %v", floppy_path, step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}

		step.Cleanup(state)

		if _, err := os.Stat(floppy_path); err == nil {
			t.Fatal("file found: %s for %v", floppy_path, step.Files)
		}
	}
}

func xxxTestStepCreateFloppy_missing(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 2
	expected := 0
	files := make([]string, count)

	prefix := "missing"

	for i := 0; i < count; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i))
	}

	lists := [][]string{
		files,
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionHalt {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); !ok {
			t.Fatalf("state should not be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path")

		if floppy_path != nil {
			t.Fatalf("floppy_path is not nil for %v", step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}
	}
}

func xxxTestStepCreateFloppy_notfound(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 2
	expected := 0
	files := make([]string, count)

	prefix := "notfound"

	for i := 0; i < count; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i))
	}

	lists := [][]string{
		[]string{dir + string(os.PathSeparator) + prefix + "*"},
		[]string{dir + string(os.PathSeparator) + prefix + "?"},
		[]string{dir + string(os.PathSeparator) + prefix + "[0123456789]"},
		[]string{dir + string(os.PathSeparator) + prefix + "[0-9]"},
		[]string{dir + string(os.PathSeparator)},
		[]string{dir},
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionContinue {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); ok {
			t.Fatalf("state should be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path").(string)

		if _, err := os.Stat(floppy_path); err != nil {
			t.Fatal("file not found: %s for %v", floppy_path, step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}

		step.Cleanup(state)

		if _, err := os.Stat(floppy_path); err == nil {
			t.Fatal("file found: %s for %v", floppy_path, step.Files)
		}
	}
}
<commit_msg>common: fix formatting directives in tests<commit_after>package common

import (
	"bytes"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"io/ioutil"
	"os"
	"path"
	"strconv"
	"testing"
)

func TestStepCreateFloppy_Impl(t *testing.T) {
	var raw interface{}
	raw = new(StepCreateFloppy)
	if _, ok := raw.(multistep.Step); !ok {
		t.Fatalf("StepCreateFloppy should be a step")
	}
}

func testStepCreateFloppyState(t *testing.T) multistep.StateBag {
	state := new(multistep.BasicStateBag)
	state.Put("ui", &packer.BasicUi{
		Reader: new(bytes.Buffer),
		Writer: new(bytes.Buffer),
	})
	return state
}

func TestStepCreateFloppy(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 10
	expected := count
	files := make([]string, count)

	prefix := "exists"
	ext := ".tmp"

	for i := 0; i < expected; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i)+ext)

		_, err := os.Create(files[i])
		if err != nil {
			t.Fatalf("err: %s", err)
		}
	}

	lists := [][]string{
		files,
		[]string{dir + string(os.PathSeparator) + prefix + "*" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "?" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "[0123456789]" + ext},
		[]string{dir + string(os.PathSeparator) + prefix + "[0-9]" + ext},
		[]string{dir + string(os.PathSeparator)},
		[]string{dir},
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionContinue {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); ok {
			t.Fatalf("state should be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path").(string)

		if _, err := os.Stat(floppy_path); err != nil {
			t.Fatalf("file not found: %s for %v", floppy_path, step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}

		step.Cleanup(state)

		if _, err := os.Stat(floppy_path); err == nil {
			t.Fatalf("file found: %s for %v", floppy_path, step.Files)
		}
	}
}

func xxxTestStepCreateFloppy_missing(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 2
	expected := 0
	files := make([]string, count)

	prefix := "missing"

	for i := 0; i < count; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i))
	}

	lists := [][]string{
		files,
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionHalt {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); !ok {
			t.Fatalf("state should not be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path")

		if floppy_path != nil {
			t.Fatalf("floppy_path is not nil for %v", step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}
	}
}

func xxxTestStepCreateFloppy_notfound(t *testing.T) {
	state := testStepCreateFloppyState(t)
	step := new(StepCreateFloppy)

	dir, err := ioutil.TempDir("", "packer")
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	defer os.RemoveAll(dir)

	count := 2
	expected := 0
	files := make([]string, count)

	prefix := "notfound"

	for i := 0; i < count; i++ {
		files[i] = path.Join(dir, prefix+strconv.Itoa(i))
	}

	lists := [][]string{
		[]string{dir + string(os.PathSeparator) + prefix + "*"},
		[]string{dir + string(os.PathSeparator) + prefix + "?"},
		[]string{dir + string(os.PathSeparator) + prefix + "[0123456789]"},
		[]string{dir + string(os.PathSeparator) + prefix + "[0-9]"},
		[]string{dir + string(os.PathSeparator)},
		[]string{dir},
	}

	for _, step.Files = range lists {
		if action := step.Run(state); action != multistep.ActionContinue {
			t.Fatalf("bad action: %#v for %v", action, step.Files)
		}

		if _, ok := state.GetOk("error"); ok {
			t.Fatalf("state should be ok for %v", step.Files)
		}

		floppy_path := state.Get("floppy_path").(string)

		if _, err := os.Stat(floppy_path); err != nil {
			t.Fatalf("file not found: %s for %v", floppy_path, step.Files)
		}

		if len(step.FilesAdded) != expected {
			t.Fatalf("expected %d, found %d for %v", expected, len(step.FilesAdded), step.Files)
		}

		step.Cleanup(state)

		if _, err := os.Stat(floppy_path); err == nil {
			t.Fatalf("file found: %s for %v", floppy_path, step.Files)
		}
	}
}
<|endoftext|>


<commit_before>package communicator

import (
	"context"
	"errors"
	"io"
	"net"
	"testing"
	"time"

	"github.com/hashicorp/terraform/terraform"
)

func TestCommunicator_new(t *testing.T) {
	r := &terraform.InstanceState{
		Ephemeral: terraform.EphemeralState{
			ConnInfo: map[string]string{
				"type": "telnet",
			},
		},
	}
	if _, err := New(r); err == nil {
		t.Fatalf("expected error with telnet")
	}

	r.Ephemeral.ConnInfo["type"] = "ssh"
	if _, err := New(r); err != nil {
		t.Fatalf("err: %v", err)
	}

	r.Ephemeral.ConnInfo["type"] = "winrm"
	if _, err := New(r); err != nil {
		t.Fatalf("err: %v", err)
	}
}
func TestRetryFunc(t *testing.T) {
	origMax := maxBackoffDelay
	maxBackoffDelay = time.Second
	origStart := initialBackoffDelay
	initialBackoffDelay = 10 * time.Millisecond

	defer func() {
		maxBackoffDelay = origMax
		initialBackoffDelay = origStart
	}()

	// succeed on the third try
	errs := []error{io.EOF, &net.OpError{Err: errors.New("ERROR")}, nil}
	count := 0

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	err := Retry(ctx, func() error {
		if count >= len(errs) {
			return errors.New("failed to stop after nil error")
		}

		err := errs[count]
		count++

		return err
	})

	if count != 3 {
		t.Fatal("retry func should have been called 3 times")
	}

	if err != nil {
		t.Fatal(err)
	}
}

func TestRetryFuncBackoff(t *testing.T) {
	origMax := maxBackoffDelay
	maxBackoffDelay = time.Second
	origStart := initialBackoffDelay
	initialBackoffDelay = 100 * time.Millisecond

	defer func() {
		maxBackoffDelay = origMax
		initialBackoffDelay = origStart
	}()

	count := 0

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	Retry(ctx, func() error {
		count++
		return io.EOF
	})

	if count > 4 {
		t.Fatalf("retry func failed to backoff. called %d times", count)
	}
}
<commit_msg>Add missing host to communicator test<commit_after>package communicator

import (
	"context"
	"errors"
	"io"
	"net"
	"testing"
	"time"

	"github.com/hashicorp/terraform/terraform"
)

func TestCommunicator_new(t *testing.T) {
	r := &terraform.InstanceState{
		Ephemeral: terraform.EphemeralState{
			ConnInfo: map[string]string{
				"type": "telnet",
				"host": "127.0.0.1",
			},
		},
	}
	if _, err := New(r); err == nil {
		t.Fatalf("expected error with telnet")
	}

	r.Ephemeral.ConnInfo["type"] = "ssh"
	if _, err := New(r); err != nil {
		t.Fatalf("err: %v", err)
	}

	r.Ephemeral.ConnInfo["type"] = "winrm"
	if _, err := New(r); err != nil {
		t.Fatalf("err: %v", err)
	}
}
func TestRetryFunc(t *testing.T) {
	origMax := maxBackoffDelay
	maxBackoffDelay = time.Second
	origStart := initialBackoffDelay
	initialBackoffDelay = 10 * time.Millisecond

	defer func() {
		maxBackoffDelay = origMax
		initialBackoffDelay = origStart
	}()

	// succeed on the third try
	errs := []error{io.EOF, &net.OpError{Err: errors.New("ERROR")}, nil}
	count := 0

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	err := Retry(ctx, func() error {
		if count >= len(errs) {
			return errors.New("failed to stop after nil error")
		}

		err := errs[count]
		count++

		return err
	})

	if count != 3 {
		t.Fatal("retry func should have been called 3 times")
	}

	if err != nil {
		t.Fatal(err)
	}
}

func TestRetryFuncBackoff(t *testing.T) {
	origMax := maxBackoffDelay
	maxBackoffDelay = time.Second
	origStart := initialBackoffDelay
	initialBackoffDelay = 100 * time.Millisecond

	defer func() {
		maxBackoffDelay = origMax
		initialBackoffDelay = origStart
	}()

	count := 0

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	Retry(ctx, func() error {
		count++
		return io.EOF
	})

	if count > 4 {
		t.Fatalf("retry func failed to backoff. called %d times", count)
	}
}
<|endoftext|>