<commit_before>package image_ecosystem

import (
	"fmt"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	exutil "github.com/openshift/origin/test/extended/util"
	dbutil "github.com/openshift/origin/test/extended/util/db"
)

var _ = g.Describe("[image_ecosystem][mongodb][Slow] openshift mongodb replication (with petset)", func() {
	defer g.GinkgoRecover()

	const templatePath = "https://raw.githubusercontent.com/sclorg/mongodb-container/master/examples/petset/mongodb-petset-persistent.yaml"

	oc := exutil.NewCLI("mongodb-petset-replica", exutil.KubeConfigPath()).Verbose()

	g.Describe("creating from a template", func() {
		g.It(fmt.Sprintf("should process and create the %q template", templatePath), func() {
			oc.SetOutputDir(exutil.TestContext.OutputDir)

			g.By("creating persistent volumes")
			_, err := exutil.SetupHostPathVolumes(
				oc.AdminKubeREST().PersistentVolumes(),
				oc.Namespace(),
				"256Mi",
				3,
			)
			o.Expect(err).NotTo(o.HaveOccurred())

			defer func() {
				// We're removing only PVs because all other things will be removed
				// together with namespace.
				err := exutil.CleanupHostPathVolumes(oc.AdminKubeREST().PersistentVolumes(), oc.Namespace())
				if err != nil {
					fmt.Fprintf(g.GinkgoWriter, "WARNING: couldn't cleanup persistent volumes: %v", err)
				}
			}()

			g.By("creating a new app")
			o.Expect(
				oc.Run("new-app").Args(
					"-f", templatePath,
					"-p", "VOLUME_CAPACITY=256Mi",
					"-p", "MEMORY_LIMIT=512Mi",
					"-p", "MONGODB_IMAGE=centos/mongodb-32-centos7",
					"-p", "MONGODB_SERVICE_NAME=mongodb-replicaset",
				).Execute(),
			).Should(o.Succeed())

			g.By("waiting for pods to running")
			podNames, err := exutil.WaitForPods(
				oc.KubeREST().Pods(oc.Namespace()),
				exutil.ParseLabelsOrDie("name=mongodb-replicaset"),
				exutil.CheckPodIsRunningFn,
				3,
				1*time.Minute,
			)
			o.Expect(err).NotTo(o.HaveOccurred())
			o.Expect(podNames).Should(o.HaveLen(3))

			g.By("expecting that we can insert a new record on primary node")
			mongo := dbutil.NewMongoDB(podNames[0])
			replicaSet := mongo.(exutil.ReplicaSet)
			_, err = replicaSet.QueryPrimary(oc, `db.test.save({ "status" : "passed" })`)
			o.Expect(err).ShouldNot(o.HaveOccurred())

			g.By("expecting that we can read a record from all members")
			for _, podName := range podNames {
				o.Expect(readRecordFromPod(oc, podName)).To(o.Succeed())
			}

			g.By("restarting replica set")
			err = oc.Run("delete").Args("pods", "--all", "-n", oc.Namespace()).Execute()
			o.Expect(err).ShouldNot(o.HaveOccurred())

			g.By("waiting for restarting of the pods")
			podNames, err = exutil.WaitForPods(
				oc.KubeREST().Pods(oc.Namespace()),
				exutil.ParseLabelsOrDie("name=mongodb-replicaset"),
				exutil.CheckPodIsRunningFn,
				3,
				1*time.Minute,
			)
			o.Expect(err).NotTo(o.HaveOccurred())
			o.Expect(podNames).Should(o.HaveLen(3))

			g.By("expecting that we can read a record from all members after its restart")
			for _, podName := range podNames {
				o.Expect(readRecordFromPod(oc, podName)).To(o.Succeed())
			}
		})
	})

})

func readRecordFromPod(oc *exutil.CLI, podName string) error {
	// don't include _id field to output because it changes every time
	findCmd := "rs.slaveOk(); printjson(db.test.find({}, {_id: 0}).toArray())"

	fmt.Fprintf(g.GinkgoWriter, "DEBUG: reading record from the pod %v\n", podName)

	mongoPod := dbutil.NewMongoDB(podName)
	// pod is running but we need to wait when it will be really ready
	// (will become a member of replica set and will finish data sync)
	return exutil.WaitForQueryOutputContains(oc, mongoPod, 1*time.Minute, false, findCmd, `{ "status" : "passed" }`)
}
<commit_msg>MongoDB replica petset test: increase time of waiting for pods.<commit_after>package image_ecosystem

import (
	"fmt"
	"time"

	g "github.com/onsi/ginkgo"
	o "github.com/onsi/gomega"

	exutil "github.com/openshift/origin/test/extended/util"
	dbutil "github.com/openshift/origin/test/extended/util/db"
)

var _ = g.Describe("[image_ecosystem][mongodb][Slow] openshift mongodb replication (with petset)", func() {
	defer g.GinkgoRecover()

	const templatePath = "https://raw.githubusercontent.com/sclorg/mongodb-container/master/examples/petset/mongodb-petset-persistent.yaml"

	oc := exutil.NewCLI("mongodb-petset-replica", exutil.KubeConfigPath()).Verbose()

	g.Describe("creating from a template", func() {
		g.It(fmt.Sprintf("should process and create the %q template", templatePath), func() {
			oc.SetOutputDir(exutil.TestContext.OutputDir)

			g.By("creating persistent volumes")
			_, err := exutil.SetupHostPathVolumes(
				oc.AdminKubeREST().PersistentVolumes(),
				oc.Namespace(),
				"256Mi",
				3,
			)
			o.Expect(err).NotTo(o.HaveOccurred())

			defer func() {
				// We're removing only PVs because all other things will be removed
				// together with namespace.
				err := exutil.CleanupHostPathVolumes(oc.AdminKubeREST().PersistentVolumes(), oc.Namespace())
				if err != nil {
					fmt.Fprintf(g.GinkgoWriter, "WARNING: couldn't cleanup persistent volumes: %v", err)
				}
			}()

			g.By("creating a new app")
			o.Expect(
				oc.Run("new-app").Args(
					"-f", templatePath,
					"-p", "VOLUME_CAPACITY=256Mi",
					"-p", "MEMORY_LIMIT=512Mi",
					"-p", "MONGODB_IMAGE=centos/mongodb-32-centos7",
					"-p", "MONGODB_SERVICE_NAME=mongodb-replicaset",
				).Execute(),
			).Should(o.Succeed())

			g.By("waiting for pods to running")
			podNames, err := exutil.WaitForPods(
				oc.KubeREST().Pods(oc.Namespace()),
				exutil.ParseLabelsOrDie("name=mongodb-replicaset"),
				exutil.CheckPodIsRunningFn,
				3,
				2*time.Minute,
			)
			o.Expect(err).NotTo(o.HaveOccurred())
			o.Expect(podNames).Should(o.HaveLen(3))

			g.By("expecting that we can insert a new record on primary node")
			mongo := dbutil.NewMongoDB(podNames[0])
			replicaSet := mongo.(exutil.ReplicaSet)
			_, err = replicaSet.QueryPrimary(oc, `db.test.save({ "status" : "passed" })`)
			o.Expect(err).ShouldNot(o.HaveOccurred())

			g.By("expecting that we can read a record from all members")
			for _, podName := range podNames {
				o.Expect(readRecordFromPod(oc, podName)).To(o.Succeed())
			}

			g.By("restarting replica set")
			err = oc.Run("delete").Args("pods", "--all", "-n", oc.Namespace()).Execute()
			o.Expect(err).ShouldNot(o.HaveOccurred())

			g.By("waiting for restarting of the pods")
			podNames, err = exutil.WaitForPods(
				oc.KubeREST().Pods(oc.Namespace()),
				exutil.ParseLabelsOrDie("name=mongodb-replicaset"),
				exutil.CheckPodIsRunningFn,
				3,
				2*time.Minute,
			)
			o.Expect(err).NotTo(o.HaveOccurred())
			o.Expect(podNames).Should(o.HaveLen(3))

			g.By("expecting that we can read a record from all members after its restart")
			for _, podName := range podNames {
				o.Expect(readRecordFromPod(oc, podName)).To(o.Succeed())
			}
		})
	})

})

func readRecordFromPod(oc *exutil.CLI, podName string) error {
	// don't include _id field to output because it changes every time
	findCmd := "rs.slaveOk(); printjson(db.test.find({}, {_id: 0}).toArray())"

	fmt.Fprintf(g.GinkgoWriter, "DEBUG: reading record from the pod %v\n", podName)

	mongoPod := dbutil.NewMongoDB(podName)
	// pod is running but we need to wait when it will be really ready
	// (will become a member of replica set and will finish data sync)
	return exutil.WaitForQueryOutputContains(oc, mongoPod, 1*time.Minute, false, findCmd, `{ "status" : "passed" }`)
}
<|endoftext|>


<commit_before>/*
Copyright 2015 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 benchmark

import (
	"fmt"
	"testing"
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/kubernetes/pkg/kubelet/apis"
	"k8s.io/kubernetes/test/integration/framework"
	testutils "k8s.io/kubernetes/test/utils"

	"k8s.io/klog"
)

var (
	defaultNodeStrategy = &testutils.TrivialNodePrepareStrategy{}
)

// BenchmarkScheduling benchmarks the scheduling rate when the cluster has
// various quantities of nodes and scheduled pods.
func BenchmarkScheduling(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 100, existingPods: 0, minPods: 100},
		{nodes: 100, existingPods: 1000, minPods: 100},
		{nodes: 1000, existingPods: 0, minPods: 100},
		{nodes: 1000, existingPods: 1000, minPods: 100},
	}
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("rc1")
	testStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("rc2")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, defaultNodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingPodAntiAffinity benchmarks the scheduling rate of pods with
// PodAntiAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingPodAntiAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithPodAntiAffinity(
		map[string]string{"name": "test", "color": "green"},
		map[string]string{"color": "green"})
	// The test strategy creates pods with anti-affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, defaultNodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingPodAffinity benchmarks the scheduling rate of pods with
// PodAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingPodAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithPodAffinity(
		map[string]string{"foo": ""},
		map[string]string{"foo": ""},
	)
	// The test strategy creates pods with affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	nodeStrategy := testutils.NewLabelNodePrepareStrategy(apis.LabelZoneFailureDomain, "zone1")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, nodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingNodeAffinity benchmarks the scheduling rate of pods with
// NodeAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingNodeAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithNodeAffinity(apis.LabelZoneFailureDomain, []string{"zone1", "zone2"})
	// The test strategy creates pods with node-affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	nodeStrategy := testutils.NewLabelNodePrepareStrategy(apis.LabelZoneFailureDomain, "zone1")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, nodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// makeBasePodWithPodAntiAffinity creates a Pod object to be used as a template.
// The Pod has a PodAntiAffinity requirement against pods with the given labels.
func makeBasePodWithPodAntiAffinity(podLabels, affinityLabels map[string]string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "anit-affinity-pod-",
			Labels:       podLabels,
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		PodAntiAffinity: &v1.PodAntiAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{
				{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: affinityLabels,
					},
					TopologyKey: apis.LabelHostname,
				},
			},
		},
	}
	return basePod
}

// makeBasePodWithPodAffinity creates a Pod object to be used as a template.
// The Pod has a PodAffinity requirement against pods with the given labels.
func makeBasePodWithPodAffinity(podLabels, affinityZoneLabels map[string]string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "affinity-pod-",
			Labels:       podLabels,
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		PodAffinity: &v1.PodAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{
				{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: affinityZoneLabels,
					},
					TopologyKey: apis.LabelZoneFailureDomain,
				},
			},
		},
	}
	return basePod
}

// makeBasePodWithNodeAffinity creates a Pod object to be used as a template.
// The Pod has a NodeAffinity requirement against nodes with the given expressions.
func makeBasePodWithNodeAffinity(key string, vals []string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "node-affinity-",
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		NodeAffinity: &v1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
				NodeSelectorTerms: []v1.NodeSelectorTerm{
					{
						MatchExpressions: []v1.NodeSelectorRequirement{
							{
								Key:      key,
								Operator: v1.NodeSelectorOpIn,
								Values:   vals,
							},
						},
					},
				},
			},
		},
	}
	return basePod
}

// benchmarkScheduling benchmarks scheduling rate with specific number of nodes
// and specific number of pods already scheduled.
// This will schedule numExistingPods pods before the benchmark starts, and at
// least minPods pods during the benchmark.
func benchmarkScheduling(numNodes, numExistingPods, minPods int,
	nodeStrategy testutils.PrepareNodeStrategy,
	setupPodStrategy, testPodStrategy testutils.TestPodCreateStrategy,
	b *testing.B) {
	if b.N < minPods {
		b.N = minPods
	}
	schedulerConfigFactory, finalFunc := mustSetupScheduler()
	defer finalFunc()
	c := schedulerConfigFactory.GetClient()

	nodePreparer := framework.NewIntegrationTestNodePreparer(
		c,
		[]testutils.CountToStrategy{{Count: numNodes, Strategy: nodeStrategy}},
		"scheduler-perf-",
	)
	if err := nodePreparer.PrepareNodes(); err != nil {
		klog.Fatalf("%v", err)
	}
	defer nodePreparer.CleanupNodes()

	config := testutils.NewTestPodCreatorConfig()
	config.AddStrategy("sched-test", numExistingPods, setupPodStrategy)
	podCreator := testutils.NewTestPodCreator(c, config)
	podCreator.CreatePods()

	for {
		scheduled, err := schedulerConfigFactory.GetScheduledPodLister().List(labels.Everything())
		if err != nil {
			klog.Fatalf("%v", err)
		}
		if len(scheduled) >= numExistingPods {
			break
		}
		time.Sleep(1 * time.Second)
	}
	// start benchmark
	b.ResetTimer()
	config = testutils.NewTestPodCreatorConfig()
	config.AddStrategy("sched-test", b.N, testPodStrategy)
	podCreator = testutils.NewTestPodCreator(c, config)
	podCreator.CreatePods()
	for {
		// This can potentially affect performance of scheduler, since List() is done under mutex.
		// TODO: Setup watch on apiserver and wait until all pods scheduled.
		scheduled, err := schedulerConfigFactory.GetScheduledPodLister().List(labels.Everything())
		if err != nil {
			klog.Fatalf("%v", err)
		}
		if len(scheduled) >= numExistingPods+b.N {
			break
		}
		// Note: This might introduce slight deviation in accuracy of benchmark results.
		// Since the total amount of time is relatively large, it might not be a concern.
		time.Sleep(100 * time.Millisecond)
	}
}
<commit_msg>adding 5000 nodes test to scheduler benchmark<commit_after>/*
Copyright 2015 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 benchmark

import (
	"fmt"
	"testing"
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/kubernetes/pkg/kubelet/apis"
	"k8s.io/kubernetes/test/integration/framework"
	testutils "k8s.io/kubernetes/test/utils"

	"k8s.io/klog"
)

var (
	defaultNodeStrategy = &testutils.TrivialNodePrepareStrategy{}
)

// BenchmarkScheduling benchmarks the scheduling rate when the cluster has
// various quantities of nodes and scheduled pods.
func BenchmarkScheduling(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 100, existingPods: 0, minPods: 100},
		{nodes: 100, existingPods: 1000, minPods: 100},
		{nodes: 1000, existingPods: 0, minPods: 100},
		{nodes: 1000, existingPods: 1000, minPods: 100},
		{nodes: 5000, existingPods: 1000, minPods: 1000},
	}
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("rc1")
	testStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("rc2")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, defaultNodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingPodAntiAffinity benchmarks the scheduling rate of pods with
// PodAntiAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingPodAntiAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
		{nodes: 5000, existingPods: 1000, minPods: 1000},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithPodAntiAffinity(
		map[string]string{"name": "test", "color": "green"},
		map[string]string{"color": "green"})
	// The test strategy creates pods with anti-affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, defaultNodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingPodAffinity benchmarks the scheduling rate of pods with
// PodAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingPodAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
		{nodes: 5000, existingPods: 1000, minPods: 1000},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithPodAffinity(
		map[string]string{"foo": ""},
		map[string]string{"foo": ""},
	)
	// The test strategy creates pods with affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	nodeStrategy := testutils.NewLabelNodePrepareStrategy(apis.LabelZoneFailureDomain, "zone1")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, nodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// BenchmarkSchedulingNodeAffinity benchmarks the scheduling rate of pods with
// NodeAffinity rules when the cluster has various quantities of nodes and
// scheduled pods.
func BenchmarkSchedulingNodeAffinity(b *testing.B) {
	tests := []struct{ nodes, existingPods, minPods int }{
		{nodes: 500, existingPods: 250, minPods: 250},
		{nodes: 500, existingPods: 5000, minPods: 250},
		{nodes: 1000, existingPods: 1000, minPods: 500},
		{nodes: 5000, existingPods: 1000, minPods: 1000},
	}
	// The setup strategy creates pods with no affinity rules.
	setupStrategy := testutils.NewSimpleWithControllerCreatePodStrategy("setup")
	testBasePod := makeBasePodWithNodeAffinity(apis.LabelZoneFailureDomain, []string{"zone1", "zone2"})
	// The test strategy creates pods with node-affinity for each other.
	testStrategy := testutils.NewCustomCreatePodStrategy(testBasePod)
	nodeStrategy := testutils.NewLabelNodePrepareStrategy(apis.LabelZoneFailureDomain, "zone1")
	for _, test := range tests {
		name := fmt.Sprintf("%vNodes/%vPods", test.nodes, test.existingPods)
		b.Run(name, func(b *testing.B) {
			benchmarkScheduling(test.nodes, test.existingPods, test.minPods, nodeStrategy, setupStrategy, testStrategy, b)
		})
	}
}

// makeBasePodWithPodAntiAffinity creates a Pod object to be used as a template.
// The Pod has a PodAntiAffinity requirement against pods with the given labels.
func makeBasePodWithPodAntiAffinity(podLabels, affinityLabels map[string]string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "anit-affinity-pod-",
			Labels:       podLabels,
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		PodAntiAffinity: &v1.PodAntiAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{
				{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: affinityLabels,
					},
					TopologyKey: apis.LabelHostname,
				},
			},
		},
	}
	return basePod
}

// makeBasePodWithPodAffinity creates a Pod object to be used as a template.
// The Pod has a PodAffinity requirement against pods with the given labels.
func makeBasePodWithPodAffinity(podLabels, affinityZoneLabels map[string]string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "affinity-pod-",
			Labels:       podLabels,
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		PodAffinity: &v1.PodAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: []v1.PodAffinityTerm{
				{
					LabelSelector: &metav1.LabelSelector{
						MatchLabels: affinityZoneLabels,
					},
					TopologyKey: apis.LabelZoneFailureDomain,
				},
			},
		},
	}
	return basePod
}

// makeBasePodWithNodeAffinity creates a Pod object to be used as a template.
// The Pod has a NodeAffinity requirement against nodes with the given expressions.
func makeBasePodWithNodeAffinity(key string, vals []string) *v1.Pod {
	basePod := &v1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			GenerateName: "node-affinity-",
		},
		Spec: testutils.MakePodSpec(),
	}
	basePod.Spec.Affinity = &v1.Affinity{
		NodeAffinity: &v1.NodeAffinity{
			RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
				NodeSelectorTerms: []v1.NodeSelectorTerm{
					{
						MatchExpressions: []v1.NodeSelectorRequirement{
							{
								Key:      key,
								Operator: v1.NodeSelectorOpIn,
								Values:   vals,
							},
						},
					},
				},
			},
		},
	}
	return basePod
}

// benchmarkScheduling benchmarks scheduling rate with specific number of nodes
// and specific number of pods already scheduled.
// This will schedule numExistingPods pods before the benchmark starts, and at
// least minPods pods during the benchmark.
func benchmarkScheduling(numNodes, numExistingPods, minPods int,
	nodeStrategy testutils.PrepareNodeStrategy,
	setupPodStrategy, testPodStrategy testutils.TestPodCreateStrategy,
	b *testing.B) {
	if b.N < minPods {
		b.N = minPods
	}
	schedulerConfigFactory, finalFunc := mustSetupScheduler()
	defer finalFunc()
	c := schedulerConfigFactory.GetClient()

	nodePreparer := framework.NewIntegrationTestNodePreparer(
		c,
		[]testutils.CountToStrategy{{Count: numNodes, Strategy: nodeStrategy}},
		"scheduler-perf-",
	)
	if err := nodePreparer.PrepareNodes(); err != nil {
		klog.Fatalf("%v", err)
	}
	defer nodePreparer.CleanupNodes()

	config := testutils.NewTestPodCreatorConfig()
	config.AddStrategy("sched-test", numExistingPods, setupPodStrategy)
	podCreator := testutils.NewTestPodCreator(c, config)
	podCreator.CreatePods()

	for {
		scheduled, err := schedulerConfigFactory.GetScheduledPodLister().List(labels.Everything())
		if err != nil {
			klog.Fatalf("%v", err)
		}
		if len(scheduled) >= numExistingPods {
			break
		}
		time.Sleep(1 * time.Second)
	}
	// start benchmark
	b.ResetTimer()
	config = testutils.NewTestPodCreatorConfig()
	config.AddStrategy("sched-test", b.N, testPodStrategy)
	podCreator = testutils.NewTestPodCreator(c, config)
	podCreator.CreatePods()
	for {
		// This can potentially affect performance of scheduler, since List() is done under mutex.
		// TODO: Setup watch on apiserver and wait until all pods scheduled.
		scheduled, err := schedulerConfigFactory.GetScheduledPodLister().List(labels.Everything())
		if err != nil {
			klog.Fatalf("%v", err)
		}
		if len(scheduled) >= numExistingPods+b.N {
			break
		}
		// Note: This might introduce slight deviation in accuracy of benchmark results.
		// Since the total amount of time is relatively large, it might not be a concern.
		time.Sleep(100 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package filestorage_test

// The filestorage structs are used as stubs in tests.
// The tests defined herein are simple smoke tests for the
// required reader and writer functionality.

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/storage"
	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju/osenv"
	jc "launchpad.net/juju-core/testing/checkers"
)

func TestPackage(t *testing.T) {
	gc.TestingT(t)
}

type filestorageSuite struct {
	dir    string
	reader storage.StorageReader
	writer storage.StorageWriter
}

var _ = gc.Suite(&filestorageSuite{})

func (s *filestorageSuite) SetUpTest(c *gc.C) {
	s.dir = c.MkDir()
	var err error
	s.reader, err = filestorage.NewFileStorageReader(s.dir)
	c.Assert(err, gc.IsNil)
	s.writer, err = filestorage.NewFileStorageWriter(s.dir, filestorage.UseDefaultTmpDir)
	c.Assert(err, gc.IsNil)
}

func (s *filestorageSuite) createFile(c *gc.C, name string) (fullpath string, data []byte) {
	fullpath = filepath.Join(s.dir, name)
	dir := filepath.Dir(fullpath)
	c.Assert(os.MkdirAll(dir, 0755), gc.IsNil)
	data = []byte{1, 2, 3, 4, 5}
	err := ioutil.WriteFile(fullpath, data, 0644)
	c.Assert(err, gc.IsNil)
	return fullpath, data
}

func (s *filestorageSuite) TestList(c *gc.C) {
	names := []string{
		"a/b/c",
		"a/bb",
		"a/c",
		"aa",
		"b/c/d",
	}
	for _, name := range names {
		s.createFile(c, name)
	}
	type test struct {
		prefix   string
		expected []string
	}
	for i, test := range []test{
		{"a", []string{"a/b/c", "a/bb", "a/c", "aa"}},
		{"a/b", []string{"a/b/c", "a/bb"}},
		{"a/b/c", []string{"a/b/c"}},
		{"", names},
	} {
		c.Logf("test %d: prefix=%q", i, test.prefix)
		files, err := storage.List(s.reader, test.prefix)
		c.Assert(err, gc.IsNil)
		c.Assert(files, gc.DeepEquals, test.expected)
	}
}

func (s *filestorageSuite) TestURL(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	url, err := s.reader.URL(file)
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+expectedpath)
}

func (s *filestorageSuite) TestGet(c *gc.C) {
	expectedpath, data := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	rc, err := storage.Get(s.reader, file)
	c.Assert(err, gc.IsNil)
	defer rc.Close()
	c.Assert(err, gc.IsNil)
	b, err := ioutil.ReadAll(rc)
	c.Assert(err, gc.IsNil)
	c.Assert(b, gc.DeepEquals, data)

	// Get on a non-existant path returns NotFoundError
	_, err = s.reader.Get("nowhere")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)

	// Get on a directory returns NotFoundError
	s.createFile(c, "dir/file")
	_, err = s.reader.Get("dir")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)
}

func (s *filestorageSuite) TestPut(c *gc.C) {
	data := []byte{1, 2, 3, 4, 5}
	err := s.writer.Put("test-write", bytes.NewReader(data), int64(len(data)))
	c.Assert(err, gc.IsNil)
	b, err := ioutil.ReadFile(filepath.Join(s.dir, "test-write"))
	c.Assert(err, gc.IsNil)
	c.Assert(b, gc.DeepEquals, data)
}

func (s *filestorageSuite) TestRemove(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	err := s.writer.Remove(file)
	c.Assert(err, gc.IsNil)
	_, err = ioutil.ReadFile(expectedpath)
	c.Assert(err, gc.Not(gc.IsNil))
}

func (s *filestorageSuite) TestRemoveAll(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	err := s.writer.RemoveAll()
	c.Assert(err, gc.IsNil)
	_, err = ioutil.ReadFile(expectedpath)
	c.Assert(err, gc.Not(gc.IsNil))
}

func (s *filestorageSuite) TestPutTmpDir(c *gc.C) {
	// Put should create and clean up the temporary directory if
	// tmpdir==UseDefaultTmpDir.
	err := s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, jc.Satisfies, os.IsNotExist)

	// To deal with recovering from hard failure, UseDefaultTmpDir
	// doesn't care if the temporary directory already exists. It
	// still removes it, though.
	err = os.Mkdir(s.dir+".tmp", 0755)
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, jc.Satisfies, os.IsNotExist)

	// If we explicitly set the temporary directory, it must already exist.
	s.writer, err = filestorage.NewFileStorageWriter(s.dir, s.dir+".tmp")
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, jc.Satisfies, os.IsNotExist)
	err = os.Mkdir(s.dir+".tmp", 0755)
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	// Temporary directory should not have been moved.
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, gc.IsNil)
}

func (s *filestorageSuite) TestPathRelativeToHome(c *gc.C) {
	reader, err := filestorage.NewFileStorageReader("~/foo")
	c.Assert(err, gc.IsNil)
	url, err := reader.URL("")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+osenv.Home()+"/foo")
}

func (s *filestorageSuite) TestRelativePath(c *gc.C) {
	dir := c.MkDir()
	err := os.MkdirAll(filepath.Join(dir, "a", "b", "c"), os.ModePerm)
	c.Assert(err, gc.IsNil)
	cwd, err := os.Getwd()
	c.Assert(err, gc.IsNil)
	err = os.Chdir(filepath.Join(dir, "a", "b", "c"))
	c.Assert(err, gc.IsNil)
	defer func() {
		os.Chdir(cwd)
	}()
	reader, err := filestorage.NewFileStorageReader("../..")
	c.Assert(err, gc.IsNil)
	url, err := reader.URL("")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+dir+"/a")
}
<commit_msg>Fix test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package filestorage_test

// The filestorage structs are used as stubs in tests.
// The tests defined herein are simple smoke tests for the
// required reader and writer functionality.

import (
	"bytes"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs/filestorage"
	"launchpad.net/juju-core/environs/storage"
	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/juju/osenv"
	jc "launchpad.net/juju-core/testing/checkers"
)

func TestPackage(t *testing.T) {
	gc.TestingT(t)
}

type filestorageSuite struct {
	dir    string
	reader storage.StorageReader
	writer storage.StorageWriter
}

var _ = gc.Suite(&filestorageSuite{})

func (s *filestorageSuite) SetUpTest(c *gc.C) {
	s.dir = c.MkDir()
	var err error
	s.reader, err = filestorage.NewFileStorageReader(s.dir)
	c.Assert(err, gc.IsNil)
	s.writer, err = filestorage.NewFileStorageWriter(s.dir, filestorage.UseDefaultTmpDir)
	c.Assert(err, gc.IsNil)
}

func (s *filestorageSuite) createFile(c *gc.C, name string) (fullpath string, data []byte) {
	fullpath = filepath.Join(s.dir, name)
	dir := filepath.Dir(fullpath)
	c.Assert(os.MkdirAll(dir, 0755), gc.IsNil)
	data = []byte{1, 2, 3, 4, 5}
	err := ioutil.WriteFile(fullpath, data, 0644)
	c.Assert(err, gc.IsNil)
	return fullpath, data
}

func (s *filestorageSuite) TestList(c *gc.C) {
	names := []string{
		"a/b/c",
		"a/bb",
		"a/c",
		"aa",
		"b/c/d",
	}
	for _, name := range names {
		s.createFile(c, name)
	}
	type test struct {
		prefix   string
		expected []string
	}
	for i, test := range []test{
		{"a", []string{"a/b/c", "a/bb", "a/c", "aa"}},
		{"a/b", []string{"a/b/c", "a/bb"}},
		{"a/b/c", []string{"a/b/c"}},
		{"", names},
	} {
		c.Logf("test %d: prefix=%q", i, test.prefix)
		files, err := storage.List(s.reader, test.prefix)
		c.Assert(err, gc.IsNil)
		c.Assert(files, gc.DeepEquals, test.expected)
	}
}

func (s *filestorageSuite) TestURL(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	url, err := s.reader.URL(file)
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+expectedpath)
}

func (s *filestorageSuite) TestGet(c *gc.C) {
	expectedpath, data := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	rc, err := storage.Get(s.reader, file)
	c.Assert(err, gc.IsNil)
	defer rc.Close()
	c.Assert(err, gc.IsNil)
	b, err := ioutil.ReadAll(rc)
	c.Assert(err, gc.IsNil)
	c.Assert(b, gc.DeepEquals, data)

	// Get on a non-existant path returns NotFoundError
	_, err = s.reader.Get("nowhere")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)

	// Get on a directory returns NotFoundError
	s.createFile(c, "dir/file")
	_, err = s.reader.Get("dir")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)
}

func (s *filestorageSuite) TestPut(c *gc.C) {
	data := []byte{1, 2, 3, 4, 5}
	err := s.writer.Put("test-write", bytes.NewReader(data), int64(len(data)))
	c.Assert(err, gc.IsNil)
	b, err := ioutil.ReadFile(filepath.Join(s.dir, "test-write"))
	c.Assert(err, gc.IsNil)
	c.Assert(b, gc.DeepEquals, data)
}

func (s *filestorageSuite) TestRemove(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	_, file := filepath.Split(expectedpath)
	err := s.writer.Remove(file)
	c.Assert(err, gc.IsNil)
	_, err = ioutil.ReadFile(expectedpath)
	c.Assert(err, gc.Not(gc.IsNil))
}

func (s *filestorageSuite) TestRemoveAll(c *gc.C) {
	expectedpath, _ := s.createFile(c, "test-file")
	err := s.writer.RemoveAll()
	c.Assert(err, gc.IsNil)
	_, err = ioutil.ReadFile(expectedpath)
	c.Assert(err, gc.Not(gc.IsNil))
}

func (s *filestorageSuite) TestPutTmpDir(c *gc.C) {
	// Put should create and clean up the temporary directory if
	// tmpdir==UseDefaultTmpDir.
	err := s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, jc.Satisfies, os.IsNotExist)

	// To deal with recovering from hard failure, UseDefaultTmpDir
	// doesn't care if the temporary directory already exists. It
	// still removes it, though.
	err = os.Mkdir(s.dir+".tmp", 0755)
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, jc.Satisfies, os.IsNotExist)

	// If we explicitly set the temporary directory, it must already exist.
	s.writer, err = filestorage.NewFileStorageWriter(s.dir, s.dir+".tmp")
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, jc.Satisfies, os.IsNotExist)
	err = os.Mkdir(s.dir+".tmp", 0755)
	c.Assert(err, gc.IsNil)
	err = s.writer.Put("test-write", bytes.NewReader(nil), 0)
	c.Assert(err, gc.IsNil)
	// Temporary directory should not have been moved.
	_, err = os.Stat(s.dir + ".tmp")
	c.Assert(err, gc.IsNil)
}

func (s *filestorageSuite) TestPathRelativeToHome(c *gc.C) {
	homeDir := osenv.Home()
	tempDir, err := ioutil.TempDir(homeDir, "")
	c.Assert(err, gc.IsNil)
	defer os.RemoveAll(tempDir)
	dirName := strings.Replace(tempDir, homeDir, "", -1)
	reader, err := filestorage.NewFileStorageReader(filepath.Join("~", dirName))
	c.Assert(err, gc.IsNil)
	url, err := reader.URL("")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+filepath.Join(homeDir, dirName))
}

func (s *filestorageSuite) TestRelativePath(c *gc.C) {
	dir := c.MkDir()
	err := os.MkdirAll(filepath.Join(dir, "a", "b", "c"), os.ModePerm)
	c.Assert(err, gc.IsNil)
	cwd, err := os.Getwd()
	c.Assert(err, gc.IsNil)
	err = os.Chdir(filepath.Join(dir, "a", "b", "c"))
	c.Assert(err, gc.IsNil)
	defer os.Chdir(cwd)
	reader, err := filestorage.NewFileStorageReader("../..")
	c.Assert(err, gc.IsNil)
	url, err := reader.URL("")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "file://"+dir+"/a")
}
<|endoftext|>


<commit_before>//go:build darwin || (linux && amd64)
// +build darwin linux,amd64

package snapmeta

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"os"
	"testing"
)

var (
	key = "mykey"
	val = []byte("myval")
)

func TestStoreLoadDelete(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	defer os.RemoveAll(repoPath)

	kpl := initKPL(t, repoPath)
	defer kpl.Cleanup()

	kpl.testStoreLoad(ctx, t, key, val)
	kpl.testDelete(ctx, t, key)
}

func TestConcurrency(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	defer os.RemoveAll(repoPath)

	kpl := initKPL(t, repoPath)
	defer kpl.Cleanup()

	keys := []string{"key1", "key2", "key3"}
	vals := [][]byte{[]byte("val1"), []byte("val2"), []byte("val3")}

	t.Run("storeLoad", func(t *testing.T) {
		for i := 0; i < 9; i++ {
			j := i
			t.Run(fmt.Sprint(i), func(t *testing.T) {
				t.Parallel()
				kpl.testStoreLoad(ctx, t, keys[j%3], vals[j%3])
			})
		}
	})

	t.Run("delete", func(t *testing.T) {
		for i := 0; i < 9; i++ {
			j := i
			t.Run(fmt.Sprint(i), func(t *testing.T) {
				t.Parallel()
				kpl.testDelete(ctx, t, keys[j%3])
			})
		}
	})
}

// Store and test that subsequent Load succeeds.
func (kpl *KopiaPersisterLight) testStoreLoad(ctx context.Context, t *testing.T, key string, val []byte) { //nolint:thelper
	err := kpl.Store(ctx, key, val)
	assertNoError(t, err)

	valOut, err := kpl.Load(ctx, key)
	assertNoError(t, err)

	if !bytes.Equal(valOut, val) {
		t.Fatal("loaded value does not equal stored value", valOut, val)
	}
}

// Delete and test that subsequent Load fails.
func (kpl *KopiaPersisterLight) testDelete(ctx context.Context, t *testing.T, key string) { //nolint:thelper
	kpl.Delete(ctx, key)

	_, err := kpl.Load(ctx, key)
	log.Println("err:", err)

	if err == nil {
		t.Fatal("snapshot was not deleted properly")
	}
}

func TestPersistence(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	kpl := initKPL(t, repoPath)

	// Persistence directory should be set.
	if persistDir := kpl.GetPersistDir(); persistDir == "" {
		t.Error("could not get persistence directory")
	}

	// These are no-ops and should always succeed.
	err = kpl.LoadMetadata()
	assertNoError(t, err)
	err = kpl.FlushMetadata()
	assertNoError(t, err)

	// Store and cleanup kpl
	err = kpl.Store(ctx, key, val)
	assertNoError(t, err)

	kpl.Cleanup()

	// Re-initialize and Load
	kpl = initKPL(t, repoPath)
	valOut, err := kpl.Load(ctx, key)
	assertNoError(t, err)

	if !bytes.Equal(valOut, val) {
		t.Fatal("loaded value does not equal stored value")
	}

	kpl.Cleanup()
	os.RemoveAll(repoPath)
}

func TestS3Connect(t *testing.T) {
	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	// Test the S3 code path by attempting to connect to a nonexistent bucket.
	t.Setenv(S3BucketNameEnvKey, "does-not-exist")

	kpl, err := NewPersisterLight("")
	assertNoError(t, err)

	if err := kpl.ConnectOrCreateRepo(repoPath); err == nil {
		t.Error("should not be able to connect to nonexistent S3 bucket")
	}

	kpl.Cleanup()
	os.RemoveAll(repoPath)
}

func initKPL(t *testing.T, repoPath string) *KopiaPersisterLight { //nolint:thelper
	os.Unsetenv(S3BucketNameEnvKey)

	kpl, err := NewPersisterLight("")
	assertNoError(t, err)

	err = kpl.ConnectOrCreateRepo(repoPath)
	assertNoError(t, err)

	return kpl
}

func assertNoError(t *testing.T, err error) {
	t.Helper()

	if err != nil {
		t.Errorf("err: %v", err)
	}
}
<commit_msg>skip TestS3Connect() until we fix the S3 regression (#2246)<commit_after>//go:build darwin || (linux && amd64)
// +build darwin linux,amd64

package snapmeta

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"os"
	"testing"
)

var (
	key = "mykey"
	val = []byte("myval")
)

func TestStoreLoadDelete(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	defer os.RemoveAll(repoPath)

	kpl := initKPL(t, repoPath)
	defer kpl.Cleanup()

	kpl.testStoreLoad(ctx, t, key, val)
	kpl.testDelete(ctx, t, key)
}

func TestConcurrency(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	defer os.RemoveAll(repoPath)

	kpl := initKPL(t, repoPath)
	defer kpl.Cleanup()

	keys := []string{"key1", "key2", "key3"}
	vals := [][]byte{[]byte("val1"), []byte("val2"), []byte("val3")}

	t.Run("storeLoad", func(t *testing.T) {
		for i := 0; i < 9; i++ {
			j := i
			t.Run(fmt.Sprint(i), func(t *testing.T) {
				t.Parallel()
				kpl.testStoreLoad(ctx, t, keys[j%3], vals[j%3])
			})
		}
	})

	t.Run("delete", func(t *testing.T) {
		for i := 0; i < 9; i++ {
			j := i
			t.Run(fmt.Sprint(i), func(t *testing.T) {
				t.Parallel()
				kpl.testDelete(ctx, t, keys[j%3])
			})
		}
	})
}

// Store and test that subsequent Load succeeds.
func (kpl *KopiaPersisterLight) testStoreLoad(ctx context.Context, t *testing.T, key string, val []byte) { //nolint:thelper
	err := kpl.Store(ctx, key, val)
	assertNoError(t, err)

	valOut, err := kpl.Load(ctx, key)
	assertNoError(t, err)

	if !bytes.Equal(valOut, val) {
		t.Fatal("loaded value does not equal stored value", valOut, val)
	}
}

// Delete and test that subsequent Load fails.
func (kpl *KopiaPersisterLight) testDelete(ctx context.Context, t *testing.T, key string) { //nolint:thelper
	kpl.Delete(ctx, key)

	_, err := kpl.Load(ctx, key)
	log.Println("err:", err)

	if err == nil {
		t.Fatal("snapshot was not deleted properly")
	}
}

func TestPersistence(t *testing.T) {
	ctx := context.Background()

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	kpl := initKPL(t, repoPath)

	// Persistence directory should be set.
	if persistDir := kpl.GetPersistDir(); persistDir == "" {
		t.Error("could not get persistence directory")
	}

	// These are no-ops and should always succeed.
	err = kpl.LoadMetadata()
	assertNoError(t, err)
	err = kpl.FlushMetadata()
	assertNoError(t, err)

	// Store and cleanup kpl
	err = kpl.Store(ctx, key, val)
	assertNoError(t, err)

	kpl.Cleanup()

	// Re-initialize and Load
	kpl = initKPL(t, repoPath)
	valOut, err := kpl.Load(ctx, key)
	assertNoError(t, err)

	if !bytes.Equal(valOut, val) {
		t.Fatal("loaded value does not equal stored value")
	}

	kpl.Cleanup()
	os.RemoveAll(repoPath)
}

func TestS3Connect(t *testing.T) {
	t.Skip("TODO - fix me")

	repoPath, err := os.MkdirTemp("", "kopia-test-repo-")
	assertNoError(t, err)

	// Test the S3 code path by attempting to connect to a nonexistent bucket.
	t.Setenv(S3BucketNameEnvKey, "does-not-exist")

	kpl, err := NewPersisterLight("")
	assertNoError(t, err)

	if err := kpl.ConnectOrCreateRepo(repoPath); err == nil {
		t.Error("should not be able to connect to nonexistent S3 bucket")
	}

	kpl.Cleanup()
	os.RemoveAll(repoPath)
}

func initKPL(t *testing.T, repoPath string) *KopiaPersisterLight { //nolint:thelper
	os.Unsetenv(S3BucketNameEnvKey)

	kpl, err := NewPersisterLight("")
	assertNoError(t, err)

	err = kpl.ConnectOrCreateRepo(repoPath)
	assertNoError(t, err)

	return kpl
}

func assertNoError(t *testing.T, err error) {
	t.Helper()

	if err != nil {
		t.Errorf("err: %v", err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

func GetRPCClient() (ret *rpc.Client, xp rpc.Transporter, err error) {
	return GetRPCClientWithContext(G)
}

func getSocketWithRetry(g *libkb.GlobalContext) (xp rpc.Transporter, err error) {
	return getSocket(g, true)
}

func getSocketNoRetry(g *libkb.GlobalContext) (xp rpc.Transporter, err error) {
	return getSocket(g, false)
}

func getSocket(g *libkb.GlobalContext, clearError bool) (xp rpc.Transporter, err error) {
	var isNew bool
	_, xp, isNew, err = g.GetSocket(clearError)
	if err == nil && isNew {
		introduceMyself(g, xp)
	}
	return xp, err
}

func GetRPCClientWithContext(g *libkb.GlobalContext) (ret *rpc.Client, xp rpc.Transporter, err error) {
	if xp, err = getSocketNoRetry(g); err == nil {
		ret = rpc.NewClient(xp, libkb.ErrorUnwrapper{})
	}
	return
}

func GetRPCServer(g *libkb.GlobalContext) (ret *rpc.Server, xp rpc.Transporter, err error) {
	if xp, err = getSocket(g, false); err == nil {
		ret = rpc.NewServer(xp, libkb.WrapError)
	}
	if err != nil {
		DiagnoseSocketError(g.UI, err)
	}
	return
}

func GetSignupClient(g *libkb.GlobalContext) (cli keybase1.SignupClient, err error) {
	var rpc *rpc.Client
	if rpc, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.SignupClient{Cli: rpc}
	}
	return
}

func GetConfigClient(g *libkb.GlobalContext) (cli keybase1.ConfigClient, err error) {
	var rpc *rpc.Client
	if rpc, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.ConfigClient{Cli: rpc}
	}
	return
}

func GetSaltpackClient(g *libkb.GlobalContext) (cli keybase1.SaltpackClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.SaltpackClient{Cli: rcli}
	}
	return
}

func GetLoginClient(g *libkb.GlobalContext) (cli keybase1.LoginClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.LoginClient{Cli: rcli}
	}
	return
}

func GetLogClient(g *libkb.GlobalContext) (cli keybase1.LogClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.LogClient{Cli: rcli}
	}
	return
}

func RegisterProtocolsWithContext(prots []rpc.Protocol, g *libkb.GlobalContext) (err error) {
	var srv *rpc.Server
	if srv, _, err = GetRPCServer(g); err != nil {
		return
	}
	prots = append(prots, NewLogUIProtocol())
	for _, p := range prots {
		if err = srv.Register(p); err != nil {
			if _, ok := err.(rpc.AlreadyRegisteredError); !ok {
				return err
			}
			err = nil
		}
	}
	return
}

func RegisterProtocols(prots []rpc.Protocol) (err error) {
	return RegisterProtocolsWithContext(prots, G)
}

func GetIdentifyClient(g *libkb.GlobalContext) (cli keybase1.IdentifyClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.IdentifyClient{Cli: rcli}
	}
	return
}

func GetProveClient() (cli keybase1.ProveClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.ProveClient{Cli: rcli}
	}
	return
}

func GetTrackClient(g *libkb.GlobalContext) (cli keybase1.TrackClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.TrackClient{Cli: rcli}
	}
	return
}

func GetDeviceClient() (cli keybase1.DeviceClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.DeviceClient{Cli: rcli}
	}
	return
}

func GetUserClient() (cli keybase1.UserClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.UserClient{Cli: rcli}
	}
	return
}

func GetSigsClient() (cli keybase1.SigsClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.SigsClient{Cli: rcli}
	}
	return
}

func GetPGPClient() (cli keybase1.PGPClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.PGPClient{Cli: rcli}
	}
	return
}

func GetRevokeClient() (cli keybase1.RevokeClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.RevokeClient{Cli: rcli}
	}
	return
}

func GetBTCClient(g *libkb.GlobalContext) (cli keybase1.BTCClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.BTCClient{Cli: rcli}
	}
	return
}

func GetCtlClient(g *libkb.GlobalContext) (cli keybase1.CtlClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.CtlClient{Cli: rcli}
	}
	return
}

func GetAccountClient(g *libkb.GlobalContext) (cli keybase1.AccountClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.AccountClient{Cli: rcli}
	}
	return
}

func GetFavoriteClient() (cli keybase1.FavoriteClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.FavoriteClient{Cli: rcli}
	}
	return
}

func GetNotifyCtlClient(g *libkb.GlobalContext) (cli keybase1.NotifyCtlClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.NotifyCtlClient{Cli: rcli}
	return cli, nil
}

func GetKBFSClient(g *libkb.GlobalContext) (cli keybase1.KbfsClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.KbfsClient{Cli: rcli}
	return cli, nil
}

func GetUpdateClient(g *libkb.GlobalContext) (cli keybase1.UpdateClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.UpdateClient{Cli: rcli}
	return cli, nil
}

func GetSecretKeysClient(g *libkb.GlobalContext) (cli keybase1.SecretKeysClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.SecretKeysClient{Cli: rcli}
	return cli, nil
}

func introduceMyself(g *libkb.GlobalContext, xp rpc.Transporter) error {
	cli := rpc.NewClient(xp, libkb.ErrorUnwrapper{})
	ccli := keybase1.ConfigClient{Cli: cli}
	return ccli.HelloIAm(context.TODO(), g.GetMyClientDetails())
}
<commit_msg>better style<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.

package client

import (
	"github.com/keybase/client/go/libkb"
	keybase1 "github.com/keybase/client/go/protocol"
	rpc "github.com/keybase/go-framed-msgpack-rpc"
	"golang.org/x/net/context"
)

func GetRPCClient() (ret *rpc.Client, xp rpc.Transporter, err error) {
	return GetRPCClientWithContext(G)
}

func getSocketWithRetry(g *libkb.GlobalContext) (xp rpc.Transporter, err error) {
	return getSocket(g, true)
}

func getSocketNoRetry(g *libkb.GlobalContext) (xp rpc.Transporter, err error) {
	return getSocket(g, false)
}

func getSocket(g *libkb.GlobalContext, clearError bool) (xp rpc.Transporter, err error) {
	var isNew bool
	_, xp, isNew, err = g.GetSocket(clearError)
	if err == nil && isNew {
		introduceMyself(g, xp)
	}
	return xp, err
}

func GetRPCClientWithContext(g *libkb.GlobalContext) (ret *rpc.Client, xp rpc.Transporter, err error) {
	if xp, err = getSocketNoRetry(g); err == nil {
		ret = rpc.NewClient(xp, libkb.ErrorUnwrapper{})
	}
	return
}

func GetRPCServer(g *libkb.GlobalContext) (ret *rpc.Server, xp rpc.Transporter, err error) {
	if xp, err = getSocketNoRetry(g); err == nil {
		ret = rpc.NewServer(xp, libkb.WrapError)
	}
	if err != nil {
		DiagnoseSocketError(g.UI, err)
	}
	return
}

func GetSignupClient(g *libkb.GlobalContext) (cli keybase1.SignupClient, err error) {
	var rpc *rpc.Client
	if rpc, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.SignupClient{Cli: rpc}
	}
	return
}

func GetConfigClient(g *libkb.GlobalContext) (cli keybase1.ConfigClient, err error) {
	var rpc *rpc.Client
	if rpc, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.ConfigClient{Cli: rpc}
	}
	return
}

func GetSaltpackClient(g *libkb.GlobalContext) (cli keybase1.SaltpackClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.SaltpackClient{Cli: rcli}
	}
	return
}

func GetLoginClient(g *libkb.GlobalContext) (cli keybase1.LoginClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.LoginClient{Cli: rcli}
	}
	return
}

func GetLogClient(g *libkb.GlobalContext) (cli keybase1.LogClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.LogClient{Cli: rcli}
	}
	return
}

func RegisterProtocolsWithContext(prots []rpc.Protocol, g *libkb.GlobalContext) (err error) {
	var srv *rpc.Server
	if srv, _, err = GetRPCServer(g); err != nil {
		return
	}
	prots = append(prots, NewLogUIProtocol())
	for _, p := range prots {
		if err = srv.Register(p); err != nil {
			if _, ok := err.(rpc.AlreadyRegisteredError); !ok {
				return err
			}
			err = nil
		}
	}
	return
}

func RegisterProtocols(prots []rpc.Protocol) (err error) {
	return RegisterProtocolsWithContext(prots, G)
}

func GetIdentifyClient(g *libkb.GlobalContext) (cli keybase1.IdentifyClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.IdentifyClient{Cli: rcli}
	}
	return
}

func GetProveClient() (cli keybase1.ProveClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.ProveClient{Cli: rcli}
	}
	return
}

func GetTrackClient(g *libkb.GlobalContext) (cli keybase1.TrackClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.TrackClient{Cli: rcli}
	}
	return
}

func GetDeviceClient() (cli keybase1.DeviceClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.DeviceClient{Cli: rcli}
	}
	return
}

func GetUserClient() (cli keybase1.UserClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.UserClient{Cli: rcli}
	}
	return
}

func GetSigsClient() (cli keybase1.SigsClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.SigsClient{Cli: rcli}
	}
	return
}

func GetPGPClient() (cli keybase1.PGPClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.PGPClient{Cli: rcli}
	}
	return
}

func GetRevokeClient() (cli keybase1.RevokeClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.RevokeClient{Cli: rcli}
	}
	return
}

func GetBTCClient(g *libkb.GlobalContext) (cli keybase1.BTCClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.BTCClient{Cli: rcli}
	}
	return
}

func GetCtlClient(g *libkb.GlobalContext) (cli keybase1.CtlClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.CtlClient{Cli: rcli}
	}
	return
}

func GetAccountClient(g *libkb.GlobalContext) (cli keybase1.AccountClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClientWithContext(g); err == nil {
		cli = keybase1.AccountClient{Cli: rcli}
	}
	return
}

func GetFavoriteClient() (cli keybase1.FavoriteClient, err error) {
	var rcli *rpc.Client
	if rcli, _, err = GetRPCClient(); err == nil {
		cli = keybase1.FavoriteClient{Cli: rcli}
	}
	return
}

func GetNotifyCtlClient(g *libkb.GlobalContext) (cli keybase1.NotifyCtlClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.NotifyCtlClient{Cli: rcli}
	return cli, nil
}

func GetKBFSClient(g *libkb.GlobalContext) (cli keybase1.KbfsClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.KbfsClient{Cli: rcli}
	return cli, nil
}

func GetUpdateClient(g *libkb.GlobalContext) (cli keybase1.UpdateClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.UpdateClient{Cli: rcli}
	return cli, nil
}

func GetSecretKeysClient(g *libkb.GlobalContext) (cli keybase1.SecretKeysClient, err error) {
	rcli, _, err := GetRPCClientWithContext(g)
	if err != nil {
		return cli, err
	}
	cli = keybase1.SecretKeysClient{Cli: rcli}
	return cli, nil
}

func introduceMyself(g *libkb.GlobalContext, xp rpc.Transporter) error {
	cli := rpc.NewClient(xp, libkb.ErrorUnwrapper{})
	ccli := keybase1.ConfigClient{Cli: cli}
	return ccli.HelloIAm(context.TODO(), g.GetMyClientDetails())
}
<|endoftext|>


<commit_before>package cmd

/*
 Copyright 2019 Crunchy Data Solutions, Inc.
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

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

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

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"strings"

	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	log "github.com/sirupsen/logrus"
)

const etcpath = "/etc/pgo/pgouser"
const pgouserenvvar = "PGOUSER"

// BasicAuthUsername and BasicAuthPassword are for BasicAuth, they are fetched from a file

var SessionCredentials msgs.BasicAuthCredentials

var caCertPool *x509.CertPool
var cert tls.Certificate
var httpclient *http.Client
var caCertPath, clientCertPath, clientKeyPath string

// StatusCheck ...
func StatusCheck(resp *http.Response) {
	log.Debugf("HTTP status code is %d", resp.StatusCode)
	if resp.StatusCode == 401 {
		fmt.Printf("Error: Authentication Failed: %d\n", resp.StatusCode)
		os.Exit(2)
	} else if resp.StatusCode != 200 {
		fmt.Printf("Error: Invalid Status Code: %d\n", resp.StatusCode)
		os.Exit(2)
	}
}

func UserHomeDir() string {
	env := "HOME"
	if runtime.GOOS == "windows" {
		env = "USERPROFILE"
	} else if runtime.GOOS == "plan9" {
		env = "home"
	}
	return os.Getenv(env)
}

func parseCredentials(dat string) msgs.BasicAuthCredentials {

	fields := strings.Split(strings.TrimSpace(dat), ":")
	log.Debugf("%v", fields)
	log.Debugf("username=[%s] password=[%s]", fields[0], fields[1])

	creds := msgs.BasicAuthCredentials{
		Username:     fields[0],
		Password:     fields[1],
		APIServerURL: APIServerURL,
	}
	return creds
}

func GetCredentials() {
	log.Debug("GetCredentials called")

	dir := UserHomeDir()
	fullPath := dir + "/" + ".pgouser"
	log.Debugf("looking in %s for credentials", fullPath)
	dat, err := ioutil.ReadFile(fullPath)
	found := false
	if err != nil {
		log.Debugf("%s not found", fullPath)
	} else {
		log.Debugf("%s found", fullPath)
		log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
		SessionCredentials = parseCredentials(string(dat))
		found = true

	}

	if !found {
		fullPath = etcpath
		dat, err = ioutil.ReadFile(fullPath)
		if err != nil {
			log.Debugf("%s not found", etcpath)
		} else {
			log.Debugf("%s found", fullPath)
			log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
			SessionCredentials = parseCredentials(string(dat))
			found = true
		}
	}

	if !found {
		pgoUser := os.Getenv(pgouserenvvar)
		if pgoUser == "" {
			fmt.Printf("Error: %s environment variable not set", pgouserenvvar)
			os.Exit(2)
		}

		fullPath = pgoUser
		log.Debugf("%s environment variable is being used at %s", pgouserenvvar, fullPath)
		dat, err = ioutil.ReadFile(fullPath)
		if err != nil {
			fmt.Printf("Error: %s file not found", fullPath)
			os.Exit(2)
		}

		log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
		SessionCredentials = parseCredentials(string(dat))
	}

	if PGO_CA_CERT != "" {
		caCertPath = PGO_CA_CERT
	} else {
		caCertPath = os.Getenv("PGO_CA_CERT")
	}

	if caCertPath == "" {
		fmt.Println("Error: PGO_CA_CERT not specified")
		os.Exit(2)
	}
	caCert, err := ioutil.ReadFile(caCertPath)
	if err != nil {
		fmt.Printf("Error: %s", err)
		fmt.Println("could not read ca certificate")
		os.Exit(2)
	}
	caCertPool = x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	if PGO_CLIENT_CERT != "" {
		clientCertPath = PGO_CLIENT_CERT
	} else {
		clientCertPath = os.Getenv("PGO_CLIENT_CERT")
	}

	if clientCertPath == "" {
		fmt.Println("Error: PGO_CLIENT_CERT not specified")
		os.Exit(2)
	}

	_, err = ioutil.ReadFile(clientCertPath)
	if err != nil {
		log.Debugf("%s not found", clientCertPath)
		os.Exit(2)
	}

	if PGO_CLIENT_KEY != "" {
		clientKeyPath = PGO_CLIENT_KEY
	} else {
		clientKeyPath = os.Getenv("PGO_CLIENT_KEY")
	}

	if clientKeyPath == "" {
		fmt.Println("Error: PGO_CLIENT_KEY not specified")
		os.Exit(2)
	}

	_, err = ioutil.ReadFile(clientKeyPath)
	if err != nil {
		log.Debugf("%s not found", clientKeyPath)
		os.Exit(2)
	}
	cert, err = tls.LoadX509KeyPair(clientCertPath, clientKeyPath)
	if err != nil {
		fmt.Printf("Error loading client certificate/key: %s\n", err)
		os.Exit(2)
	}

	log.Debug("setting up httpclient with TLS")
	httpclient = &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				RootCAs:            caCertPool,
				InsecureSkipVerify: true,
				Certificates:       []tls.Certificate{cert},
				MinVersion:         tls.VersionTLS11,
			},
		},
	}

}
<commit_msg>make PGOUSER env var the first place to look for the pgouser file<commit_after>package cmd

/*
 Copyright 2019 Crunchy Data Solutions, Inc.
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

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

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

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"runtime"
	"strings"

	msgs "github.com/crunchydata/postgres-operator/apiservermsgs"
	log "github.com/sirupsen/logrus"
)

const etcpath = "/etc/pgo/pgouser"
const pgouserenvvar = "PGOUSER"

// BasicAuthUsername and BasicAuthPassword are for BasicAuth, they are fetched from a file

var SessionCredentials msgs.BasicAuthCredentials

var caCertPool *x509.CertPool
var cert tls.Certificate
var httpclient *http.Client
var caCertPath, clientCertPath, clientKeyPath string

// StatusCheck ...
func StatusCheck(resp *http.Response) {
	log.Debugf("HTTP status code is %d", resp.StatusCode)
	if resp.StatusCode == 401 {
		fmt.Printf("Error: Authentication Failed: %d\n", resp.StatusCode)
		os.Exit(2)
	} else if resp.StatusCode != 200 {
		fmt.Printf("Error: Invalid Status Code: %d\n", resp.StatusCode)
		os.Exit(2)
	}
}

func UserHomeDir() string {
	env := "HOME"
	if runtime.GOOS == "windows" {
		env = "USERPROFILE"
	} else if runtime.GOOS == "plan9" {
		env = "home"
	}
	return os.Getenv(env)
}

func parseCredentials(dat string) msgs.BasicAuthCredentials {

	fields := strings.Split(strings.TrimSpace(dat), ":")
	log.Debugf("%v", fields)
	log.Debugf("username=[%s] password=[%s]", fields[0], fields[1])

	creds := msgs.BasicAuthCredentials{
		Username:     fields[0],
		Password:     fields[1],
		APIServerURL: APIServerURL,
	}
	return creds
}

func GetCredentials() {
	log.Debug("GetCredentials called")

	found := false
	dir := UserHomeDir()
	fullPath := dir + "/" + ".pgouser"

	//look in env var for pgouser file
	pgoUser := os.Getenv(pgouserenvvar)
	if pgoUser != "" {
		fullPath = pgoUser
		log.Debugf("%s environment variable is being used at %s", pgouserenvvar, fullPath)
		dat, err := ioutil.ReadFile(fullPath)
		if err != nil {
			fmt.Printf("Error: %s file not found", fullPath)
			os.Exit(2)
		}

		log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
		SessionCredentials = parseCredentials(string(dat))
		found = true
	}

	//look in home directory for .pgouser file
	if !found {
		log.Debugf("looking in %s for credentials", fullPath)
		dat, err := ioutil.ReadFile(fullPath)
		if err != nil {
			log.Debugf("%s not found", fullPath)
		} else {
			log.Debugf("%s found", fullPath)
			log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
			SessionCredentials = parseCredentials(string(dat))
			found = true

		}
	}

	//look in etc for pgouser file
	if !found {
		fullPath = etcpath
		dat, err := ioutil.ReadFile(fullPath)
		if err != nil {
			log.Debugf("%s not found", etcpath)
		} else {
			log.Debugf("%s found", fullPath)
			log.Debugf("pgouser file found at %s contains %s", fullPath, string(dat))
			SessionCredentials = parseCredentials(string(dat))
			found = true
		}
	}

	if !found {
		fmt.Println("could not find pgouser file")
		os.Exit(2)
	}

	if PGO_CA_CERT != "" {
		caCertPath = PGO_CA_CERT
	} else {
		caCertPath = os.Getenv("PGO_CA_CERT")
	}

	if caCertPath == "" {
		fmt.Println("Error: PGO_CA_CERT not specified")
		os.Exit(2)
	}
	caCert, err := ioutil.ReadFile(caCertPath)
	if err != nil {
		fmt.Printf("Error: %s", err)
		fmt.Println("could not read ca certificate")
		os.Exit(2)
	}
	caCertPool = x509.NewCertPool()
	caCertPool.AppendCertsFromPEM(caCert)

	if PGO_CLIENT_CERT != "" {
		clientCertPath = PGO_CLIENT_CERT
	} else {
		clientCertPath = os.Getenv("PGO_CLIENT_CERT")
	}

	if clientCertPath == "" {
		fmt.Println("Error: PGO_CLIENT_CERT not specified")
		os.Exit(2)
	}

	_, err = ioutil.ReadFile(clientCertPath)
	if err != nil {
		log.Debugf("%s not found", clientCertPath)
		os.Exit(2)
	}

	if PGO_CLIENT_KEY != "" {
		clientKeyPath = PGO_CLIENT_KEY
	} else {
		clientKeyPath = os.Getenv("PGO_CLIENT_KEY")
	}

	if clientKeyPath == "" {
		fmt.Println("Error: PGO_CLIENT_KEY not specified")
		os.Exit(2)
	}

	_, err = ioutil.ReadFile(clientKeyPath)
	if err != nil {
		log.Debugf("%s not found", clientKeyPath)
		os.Exit(2)
	}
	cert, err = tls.LoadX509KeyPair(clientCertPath, clientKeyPath)
	if err != nil {
		fmt.Printf("Error loading client certificate/key: %s\n", err)
		os.Exit(2)
	}

	log.Debug("setting up httpclient with TLS")
	httpclient = &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				RootCAs:            caCertPool,
				InsecureSkipVerify: true,
				Certificates:       []tls.Certificate{cert},
				MinVersion:         tls.VersionTLS11,
			},
		},
	}

}
<|endoftext|>


<commit_before>// Copyright (c) 2016 VMware, Inc. All Rights Reserved.
//
// This product is licensed to you under the Apache License, Version 2.0 (the "License").
// You may not use this product except in compliance with the License.
//
// This product may include a number of subcomponents with separate copyright notices and
// license terms. Your use of these subcomponents is subject to the terms and conditions
// of the subcomponent's license, as noted in the LICENSE file.

package photon

import (
	"bytes"
	"encoding/json"
)

// Contains functionality for hosts API.
type HostsAPI struct {
	client *Client
}

var hostUrl string = rootUrl + "/hosts"

// Sets host's availability zone.
func (api *HostsAPI) SetAvailabilityZone(id string, availabilityZone *HostSetAvailabilityZoneOperation) (task *Task, err error) {
	body, err := json.Marshal(availabilityZone)
	if err != nil {
		return
	}

	res, err := api.client.restClient.Post(
		api.client.Endpoint+hostUrl+"/"+id+"/set_availability_zone",
		"application/json",
		bytes.NewReader(body),
		api.client.options.TokenOptions)

	if err != nil {
		return
	}

	defer res.Body.Close()
	task, err = getTask(getError(res))
	return
}

// Gets all tasks with the specified host ID, using options to filter the results.
// If options is nil, no filtering will occur.
func (api *HostsAPI) GetTasks(id string, options *TaskGetOptions) (result *TaskList, err error) {
	uri := api.client.Endpoint + hostUrl + "/" + id + "/tasks"
	if options != nil {
		uri += getQueryString(options)
	}
	res, err := api.client.restClient.GetList(api.client.Endpoint, uri, api.client.options.TokenOptions)
	if err != nil {
		return
	}

	result = &TaskList{}
	err = json.Unmarshal(res, result)
	return
}

// provision the host with the specified id
func (api *HostsAPI) Provision(id string) (task *Task, err error) {
	body := []byte{}
	res, err := api.client.restClient.Post(
		api.client.Endpoint+hostUrl+"/"+id+"/provision",
		"application/json",
		bytes.NewReader(body),
		api.client.options.TokenOptions)
	if err != nil {
		return
	}
	defer res.Body.Close()
	task, err = getTask(getError(res))
	return
}
<commit_msg>Replace endpoint /v1/hosts/{id} to /v1/infrastructure/hosts/[id}<commit_after>// Copyright (c) 2016 VMware, Inc. All Rights Reserved.
//
// This product is licensed to you under the Apache License, Version 2.0 (the "License").
// You may not use this product except in compliance with the License.
//
// This product may include a number of subcomponents with separate copyright notices and
// license terms. Your use of these subcomponents is subject to the terms and conditions
// of the subcomponent's license, as noted in the LICENSE file.

package photon

import (
	"bytes"
	"encoding/json"
)

// Contains functionality for hosts API.
type HostsAPI struct {
	client *Client
}

var hostUrl string = rootUrl + "/infrastructure/hosts"

// Sets host's availability zone.
func (api *HostsAPI) SetAvailabilityZone(id string, availabilityZone *HostSetAvailabilityZoneOperation) (task *Task, err error) {
	body, err := json.Marshal(availabilityZone)
	if err != nil {
		return
	}

	res, err := api.client.restClient.Post(
		api.client.Endpoint+hostUrl+"/"+id+"/set_availability_zone",
		"application/json",
		bytes.NewReader(body),
		api.client.options.TokenOptions)

	if err != nil {
		return
	}

	defer res.Body.Close()
	task, err = getTask(getError(res))
	return
}

// Gets all tasks with the specified host ID, using options to filter the results.
// If options is nil, no filtering will occur.
func (api *HostsAPI) GetTasks(id string, options *TaskGetOptions) (result *TaskList, err error) {
	uri := api.client.Endpoint + hostUrl + "/" + id + "/tasks"
	if options != nil {
		uri += getQueryString(options)
	}
	res, err := api.client.restClient.GetList(api.client.Endpoint, uri, api.client.options.TokenOptions)
	if err != nil {
		return
	}

	result = &TaskList{}
	err = json.Unmarshal(res, result)
	return
}

// provision the host with the specified id
func (api *HostsAPI) Provision(id string) (task *Task, err error) {
	body := []byte{}
	res, err := api.client.restClient.Post(
		api.client.Endpoint+hostUrl+"/"+id+"/provision",
		"application/json",
		bytes.NewReader(body),
		api.client.options.TokenOptions)
	if err != nil {
		return
	}
	defer res.Body.Close()
	task, err = getTask(getError(res))
	return
}
<|endoftext|>


<commit_before>package tfhelper

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"os"
	"path"
	"strings"

	"github.com/bdastur/command"
	"github.com/bdastur/templates"
)

const (
	environments_dir = "/tmp/s3backends"
	templates_dir    = "tftemplates"
	terraform_bin    = "/usr/local/bin/terraform"
)

// func checkErr(err error, message string) {
// 	if err != nil {
// 		fmt.Printf("%s\n", message)
// 	}
// }

type ClusterSpec struct {
	Account string `json: "account"`
	Region  string `json: "region"`
}

func parseTemplate(clusterSpec ClusterSpec, templateData string) string {
	tmpl, err := template.New("spec").Parse(templateData)
	if err != nil {
		panic(err)
	}
	renderedData := new(bytes.Buffer)
	err = tmpl.Execute(renderedData, clusterSpec)
	if err != nil {
		panic(err)
	}
	return renderedData.String()
}

func getProviderDefinition(clusterSpec ClusterSpec) string {
	providerTmpl := templates.GetProviderTemplate()
	strings.Trim(providerTmpl, " ")

	renderedData := parseTemplate(clusterSpec, providerTmpl)
	return renderedData
}

func getBackendDefinition(clusterSpec ClusterSpec) string {
	backendTmpl := templates.GetBackendTemplate()
	strings.Trim(backendTmpl, " ")

	renderedData := parseTemplate(clusterSpec, backendTmpl)
	return renderedData
}

func getRemoteStateDefinition(clusterSpec ClusterSpec) string {
	remoteTmpl := templates.GetTfStateRemoteTemplate()
	strings.Trim(remoteTmpl, " ")

	renderedData := parseTemplate(clusterSpec, remoteTmpl)
	return renderedData
}

func buildClusterSpec(clusterSpecString string) (error, ClusterSpec) {
	var clusterSpec ClusterSpec

	fmt.Println("cluster spec string: ", clusterSpecString)
	jsonbytes := []byte(clusterSpecString)
	err := json.Unmarshal(jsonbytes, &clusterSpec)
	if err != nil {
		fmt.Println("Failed to Unmarsh json data: ", err)
		return nil, clusterSpec
	}

	fmt.Println("Region: ", clusterSpec.Region)
	fmt.Println("Account: ", clusterSpec.Account)

	return nil, clusterSpec
}

func setupStagingFolder(region string, account string) string {
	// Make the staging folder.
	mode := os.FileMode(0744)
	if _, err := os.Stat(environments_dir); os.IsNotExist(err) {
		os.Mkdir(environments_dir, mode)
	} else {
		fmt.Printf("Directory %s exists \n", environments_dir)
	}

	// Create a subfolder
	s := []string{region, "-", account}
	envdir_name := strings.Join(s, "")
	fmt.Printf("Env subdirname: %s \n", envdir_name)

	s3_staging_folder := path.Join(environments_dir, envdir_name)
	fmt.Printf("Staging folder: %s \n", s3_staging_folder)

	//Create staging folder.
	if _, err := os.Stat(s3_staging_folder); os.IsNotExist(err) {
		fmt.Printf("Creating new folder %s", s3_staging_folder)
		os.Mkdir(s3_staging_folder, mode)
	} else {
		fmt.Printf("Directory %s exists \n", s3_staging_folder)
	}

	return s3_staging_folder
}

func createTerraformFile(stagingFolder string, tfDefinition string) error {
	// Crete a new tf definition file in the staging environment.
	filePath := path.Join(stagingFolder, "main.tf")
	fileHandle, err := os.Create(filePath)

	count, err := fileHandle.WriteString(tfDefinition)
	if err != nil {
		fmt.Println("Failed to write data to ", filePath)
		return err
	}
	fmt.Println("Bytes written: ", count)

	return err
}

func createTemplateDefinition(clusterSpec ClusterSpec, renderBackend bool) string {
	/*
	 * Render template definitions
	 */

	// Render provider.
	providerDefinition := getProviderDefinition(clusterSpec)

	// Render Remote state.
	remoteStateDefinition := getRemoteStateDefinition(clusterSpec)

	// Render Backend.
	var s []string
	if renderBackend {
		backendDefinition := getBackendDefinition(clusterSpec)
		s = []string{providerDefinition, "\n",
			backendDefinition, "\n", remoteStateDefinition}
	} else {
		s = []string{providerDefinition, "\n", remoteStateDefinition}
	}

	// Concat definitions into a single string.
	templateDefinition := strings.Join(s, "")

	return templateDefinition
}

func BootstrapEnvironment(clusterSpecString string) {
	fmt.Println("tfhelper-go")
	region := "us-west-2"
	account := "dev1"
	fmt.Printf("Account: %s, Region: %s \n", account, region)

	//Build cluster spec object.
	err, clusterSpec := buildClusterSpec(clusterSpecString)
	if err != nil {
		fmt.Println("Failed to build cluster spec!")
		return
	}
	fmt.Println("cluster Spec: ", clusterSpec)

	// Setup staging folder.
	s3_staging_folder := setupStagingFolder(clusterSpec.Region, clusterSpec.Account)
	fmt.Println("Staging fodler ready: ", s3_staging_folder)

	/*
	 * Render template definitions
	 */
	templateDefinition := createTemplateDefinition(clusterSpec, false)
	fmt.Println("Template definition: ", templateDefinition)

	//Create Terraform definition file.
	createTerraformFile(s3_staging_folder, templateDefinition)

	// Terraform init.
	out, err := command.ExecuteCommand(terraform_bin, s3_staging_folder, "init")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return
	}
	fmt.Println("out: ", out)

	// Terraform plan
	out, err = command.ExecuteCommand(terraform_bin, s3_staging_folder, "plan")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return
	}
	fmt.Println("out: ", out)

	// Terraform apply
	out, err = command.ExecuteCommand(terraform_bin, s3_staging_folder,
		"apply", "-auto-approve")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return
	}
	fmt.Println("out: ", out)

}
<commit_msg>Code cleanup<commit_after>package tfhelper

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"os"
	"path"
	"strings"

	"github.com/bdastur/command"
	"github.com/bdastur/templates"
)

const (
	environments_dir = "/tmp/s3backends"
	templates_dir    = "tftemplates"
	terraform_bin    = "/usr/local/bin/terraform"
)

// func checkErr(err error, message string) {
// 	if err != nil {
// 		fmt.Printf("%s\n", message)
// 	}
// }

type ClusterSpec struct {
	Account string `json: "account"`
	Region  string `json: "region"`
}

func parseTemplate(clusterSpec ClusterSpec, templateData string) string {
	tmpl, err := template.New("spec").Parse(templateData)
	if err != nil {
		panic(err)
	}
	renderedData := new(bytes.Buffer)
	err = tmpl.Execute(renderedData, clusterSpec)
	if err != nil {
		panic(err)
	}
	return renderedData.String()
}

func getProviderDefinition(clusterSpec ClusterSpec) string {
	providerTmpl := templates.GetProviderTemplate()
	strings.Trim(providerTmpl, " ")

	renderedData := parseTemplate(clusterSpec, providerTmpl)
	return renderedData
}

func getBackendDefinition(clusterSpec ClusterSpec) string {
	backendTmpl := templates.GetBackendTemplate()
	strings.Trim(backendTmpl, " ")

	renderedData := parseTemplate(clusterSpec, backendTmpl)
	return renderedData
}

func getRemoteStateDefinition(clusterSpec ClusterSpec) string {
	remoteTmpl := templates.GetTfStateRemoteTemplate()
	strings.Trim(remoteTmpl, " ")

	renderedData := parseTemplate(clusterSpec, remoteTmpl)
	return renderedData
}

func buildClusterSpec(clusterSpecString string) (error, ClusterSpec) {
	var clusterSpec ClusterSpec

	fmt.Println("cluster spec string: ", clusterSpecString)
	jsonbytes := []byte(clusterSpecString)
	err := json.Unmarshal(jsonbytes, &clusterSpec)
	if err != nil {
		fmt.Println("Failed to Unmarsh json data: ", err)
		return nil, clusterSpec
	}

	fmt.Println("Region: ", clusterSpec.Region)
	fmt.Println("Account: ", clusterSpec.Account)

	return nil, clusterSpec
}

func setupStagingFolder(region string, account string) string {
	// Make the staging folder.
	mode := os.FileMode(0744)
	if _, err := os.Stat(environments_dir); os.IsNotExist(err) {
		os.Mkdir(environments_dir, mode)
	} else {
		fmt.Printf("Directory %s exists \n", environments_dir)
	}

	// Create a subfolder
	s := []string{region, "-", account}
	envdir_name := strings.Join(s, "")
	fmt.Printf("Env subdirname: %s \n", envdir_name)

	s3_staging_folder := path.Join(environments_dir, envdir_name)
	fmt.Printf("Staging folder: %s \n", s3_staging_folder)

	//Create staging folder.
	if _, err := os.Stat(s3_staging_folder); os.IsNotExist(err) {
		fmt.Printf("Creating new folder %s", s3_staging_folder)
		os.Mkdir(s3_staging_folder, mode)
	} else {
		fmt.Printf("Directory %s exists \n", s3_staging_folder)
	}

	return s3_staging_folder
}

func createTerraformFile(stagingFolder string, tfDefinition string) error {
	// Crete a new tf definition file in the staging environment.
	filePath := path.Join(stagingFolder, "main.tf")
	fileHandle, err := os.Create(filePath)

	count, err := fileHandle.WriteString(tfDefinition)
	if err != nil {
		fmt.Println("Failed to write data to ", filePath)
		return err
	}
	fmt.Println("Bytes written: ", count)

	return err
}

func createTemplateDefinition(clusterSpec ClusterSpec, renderBackend bool) string {
	/*
	 * Render template definitions
	 */

	// Render provider.
	providerDefinition := getProviderDefinition(clusterSpec)

	// Render Remote state.
	remoteStateDefinition := getRemoteStateDefinition(clusterSpec)

	// Render Backend.
	var s []string
	if renderBackend {
		backendDefinition := getBackendDefinition(clusterSpec)
		s = []string{providerDefinition, "\n",
			backendDefinition, "\n", remoteStateDefinition}
	} else {
		s = []string{providerDefinition, "\n", remoteStateDefinition}
	}

	// Concat definitions into a single string.
	templateDefinition := strings.Join(s, "")

	return templateDefinition
}

func deployTerraformState(s3_staging_folder string) error {
	// Terraform init.
	out, err := command.ExecuteCommand(terraform_bin, s3_staging_folder, "init")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return err
	}
	fmt.Println("out: ", out)

	// Terraform plan
	out, err = command.ExecuteCommand(terraform_bin, s3_staging_folder, "plan")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return err
	}
	fmt.Println("out: ", out)

	// Terraform apply
	out, err = command.ExecuteCommand(terraform_bin, s3_staging_folder,
		"apply", "-auto-approve")
	if err != nil {
		fmt.Println("Error executing command: ", err)
		return err
	}
	fmt.Println("out: ", out)
	return err
}

func BootstrapEnvironment(clusterSpecString string) {
	fmt.Println("tfhelper-go")
	region := "us-west-2"
	account := "dev1"
	fmt.Printf("Account: %s, Region: %s \n", account, region)

	//Build cluster spec object.
	err, clusterSpec := buildClusterSpec(clusterSpecString)
	if err != nil {
		fmt.Println("Failed to build cluster spec!")
		return
	}
	fmt.Println("cluster Spec: ", clusterSpec)

	// Setup staging folder.
	s3_staging_folder := setupStagingFolder(clusterSpec.Region, clusterSpec.Account)
	fmt.Println("Staging fodler ready: ", s3_staging_folder)

	/*
	 * Render template definitions
	 */
	templateDefinition := createTemplateDefinition(clusterSpec, false)
	fmt.Println("Template definition: ", templateDefinition)

	//Create Terraform definition file.
	createTerraformFile(s3_staging_folder, templateDefinition)

	// Deploy terraform state.
	err = deployTerraformState(s3_staging_folder)
	if err != nil {
		fmt.Println("Failed to deploy terraform state.")
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/cenkalti/backoff"
	logrus "github.com/sirupsen/logrus"
)

type WorkerConfig struct {
	Workers  int `json:"workers,omitempty"`
	MaxTries int `json:"max_tries,omitempty"`
}

func (c *WorkerConfig) setup() *ConfigError {
	if c.Workers < 1 {
		c.Workers = 1
	}
	if c.MaxTries < 1 {
		c.MaxTries = 5
	}
	return nil
}

type Target struct {
	Bucket    string
	Object    string
	LocalPath string
	Error     error
}

func (t *Target) URL() string {
	return fmt.Sprintf("gs://%s/%s", t.Bucket, t.Object)
}

type Targets []*Target

func (targets Targets) error() error {
	messages := []string{}
	for _, t := range targets {
		if t.Error != nil {
			messages = append(messages, t.Error.Error())
		}
	}
	if len(messages) == 0 {
		return nil
	}
	return fmt.Errorf(strings.Join(messages, "\n"))
}

type TargetWorker struct {
	name     string
	targets  chan *Target
	impl     func(bucket, object, srcPath string) error
	done     bool
	maxTries int
}

func (w *TargetWorker) run() {
	for {
		flds := logrus.Fields{}
		log.Debugln("Getting a target")
		var t *Target
		select {
		case t = <-w.targets:
		default: // Do nothing to break
		}
		if t == nil {
			log.Debugln("No target found any more")
			w.done = true
			break
		}
		if t.Error != nil {
			continue
		}

		flds["target"] = t
		log.WithFields(flds).Debugf("Worker Start to %v\n", w.name)

		f := func() error {
			return w.impl(t.Bucket, t.Object, t.LocalPath)
		}

		eb := backoff.NewExponentialBackOff()
		eb.InitialInterval = 30 * time.Second
		b := backoff.WithMaxTries(eb, uint64(w.maxTries))
		err := backoff.Retry(f, b)
		flds["error"] = err
		if err != nil {
			log.WithFields(flds).Errorf("Worker Failed to %v %v\n", w.name, t.URL())
			w.done = true
			t.Error = err
			continue
		}
		log.WithFields(flds).Debugf("Worker Finished to %v\n", w.name)
	}
}

type TargetWorkers []*TargetWorker

func (ws TargetWorkers) process(targets Targets) error {
	c := make(chan *Target, len(targets))
	for _, t := range targets {
		c <- t
	}

	for _, w := range ws {
		w.targets = c
		go w.run()
	}

	for {
		time.Sleep(100 * time.Millisecond)
		if ws.done() {
			break
		}
	}

	return targets.error()
}

func (ws TargetWorkers) done() bool {
	for _, w := range ws {
		if !w.done {
			return false
		}
	}
	return true
}
<commit_msg>:bug: Don't mark done if error<commit_after>package main

import (
	"fmt"
	"strings"
	"time"

	"github.com/cenkalti/backoff"
	logrus "github.com/sirupsen/logrus"
)

type WorkerConfig struct {
	Workers  int `json:"workers,omitempty"`
	MaxTries int `json:"max_tries,omitempty"`
}

func (c *WorkerConfig) setup() *ConfigError {
	if c.Workers < 1 {
		c.Workers = 1
	}
	if c.MaxTries < 1 {
		c.MaxTries = 5
	}
	return nil
}

type Target struct {
	Bucket    string
	Object    string
	LocalPath string
	Error     error
}

func (t *Target) URL() string {
	return fmt.Sprintf("gs://%s/%s", t.Bucket, t.Object)
}

type Targets []*Target

func (targets Targets) error() error {
	messages := []string{}
	for _, t := range targets {
		if t.Error != nil {
			messages = append(messages, t.Error.Error())
		}
	}
	if len(messages) == 0 {
		return nil
	}
	return fmt.Errorf(strings.Join(messages, "\n"))
}

type TargetWorker struct {
	name     string
	targets  chan *Target
	impl     func(bucket, object, srcPath string) error
	done     bool
	maxTries int
}

func (w *TargetWorker) run() {
	for {
		flds := logrus.Fields{}
		log.Debugln("Getting a target")
		var t *Target
		select {
		case t = <-w.targets:
		default: // Do nothing to break
		}
		if t == nil {
			log.Debugln("No target found any more")
			w.done = true
			break
		}
		if t.Error != nil {
			continue
		}

		flds["target"] = t
		log.WithFields(flds).Debugf("Worker Start to %v\n", w.name)

		f := func() error {
			return w.impl(t.Bucket, t.Object, t.LocalPath)
		}

		eb := backoff.NewExponentialBackOff()
		eb.InitialInterval = 30 * time.Second
		b := backoff.WithMaxTries(eb, uint64(w.maxTries))
		err := backoff.Retry(f, b)
		flds["error"] = err
		if err != nil {
			log.WithFields(flds).Errorf("Worker Failed to %v %v\n", w.name, t.URL())
			t.Error = err
			continue
		}
		log.WithFields(flds).Debugf("Worker Finished to %v\n", w.name)
	}
}

type TargetWorkers []*TargetWorker

func (ws TargetWorkers) process(targets Targets) error {
	c := make(chan *Target, len(targets))
	for _, t := range targets {
		c <- t
	}

	for _, w := range ws {
		w.targets = c
		go w.run()
	}

	for {
		time.Sleep(100 * time.Millisecond)
		if ws.done() {
			break
		}
	}

	return targets.error()
}

func (ws TargetWorkers) done() bool {
	for _, w := range ws {
		if !w.done {
			return false
		}
	}
	return true
}
<|endoftext|>