<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package registrydisk

import (
	"io/ioutil"
	"os"
	"os/user"

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

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
)

var _ = Describe("RegistryDisk", func() {
	tmpDir, _ := ioutil.TempDir("", "registrydisktest")
	owner, err := user.Current()
	if err != nil {
		panic(err)
	}

	VerifyDiskType := func(diskExtension string) {
		vm := v1.NewMinimalVM("fake-vm")
		vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
			Type:   "RegistryDisk:v1alpha",
			Device: "disk",
			Source: v1.DiskSource{
				Name: "someimage:v1.2.3.4",
			},
			Target: v1.DiskTarget{
				Device: "vda",
			},
		})

		// create a fake disk file
		volumeMountDir := generateVMBaseDir(vm)
		err = os.MkdirAll(volumeMountDir+"/disk0", 0750)
		Expect(err).ToNot(HaveOccurred())

		filePath := volumeMountDir + "/disk0/disk-image." + diskExtension
		_, err := os.Create(filePath)

		vm, err = MapRegistryDisks(vm)
		Expect(err).ToNot(HaveOccurred())

		// verify file gets renamed by virt-handler to prevent container from
		// removing it before VM is exited
		exists, err := diskutils.FileExists(filePath)
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(false))

		// verify file rename takes place
		exists, err = diskutils.FileExists(filePath + ".virt")
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(true))

		Expect(vm.Spec.Domain.Devices.Disks[0].Type).To(Equal("file"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Target.Device).To(Equal("vda"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver.Type).To(Equal(diskExtension))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source.File).To(Equal(filePath + ".virt"))

		err = CleanupEphemeralDisks(vm)
		exists, err = diskutils.FileExists(volumeMountDir)
		Expect(err).ToNot(HaveOccurred())

		Expect(exists).To(Equal(false))
	}

	BeforeSuite(func() {
		err := SetLocalDirectory(tmpDir)
		if err != nil {
			panic(err)
		}
		SetLocalDataOwner(owner.Username)
	})

	AfterSuite(func() {
		os.RemoveAll(tmpDir)
	})

	Describe("registry-disk", func() {
		Context("verify helper functions", func() {
			It("by verifying error when no disk is present", func() {

				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				vm, err := MapRegistryDisks(vm)
				Expect(err).To(HaveOccurred())
			})

			It("by verifying mapping of qcow2 disk", func() {
				VerifyDiskType("qcow2")
			})

			It("by verifying mapping of raw disk", func() {
				VerifyDiskType("raw")
			})

			It("by verifying container generation", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vdb",
					},
				})

				containers, volumes, err := GenerateContainers(vm)
				Expect(err).ToNot(HaveOccurred())

				Expect(len(containers)).To(Equal(2))
				Expect(len(volumes)).To(Equal(2))
			})

			It("by verifying data cleanup", func() {

				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				volumeMountDir := generateVMBaseDir(vm)
				err = os.MkdirAll(volumeMountDir, 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk0", 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk1", 0755)
				Expect(err).ToNot(HaveOccurred())

				exists, err := diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(true))

				err = CleanupEphemeralDisks(vm)
				exists, err = diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(false))

			})
		})
	})
})
<commit_msg>Simplify registry-disk unit tests by utilizing testing table<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package registrydisk

import (
	"io/ioutil"
	"os"
	"os/user"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"

	v1 "kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
)

var _ = Describe("RegistryDisk", func() {
	tmpDir, _ := ioutil.TempDir("", "registrydisktest")
	owner, err := user.Current()
	if err != nil {
		panic(err)
	}

	VerifyDiskType := func(diskExtension string) {
		vm := v1.NewMinimalVM("fake-vm")
		vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
			Type:   "RegistryDisk:v1alpha",
			Device: "disk",
			Source: v1.DiskSource{
				Name: "someimage:v1.2.3.4",
			},
			Target: v1.DiskTarget{
				Device: "vda",
			},
		})

		// create a fake disk file
		volumeMountDir := generateVMBaseDir(vm)
		err = os.MkdirAll(volumeMountDir+"/disk0", 0750)
		Expect(err).ToNot(HaveOccurred())

		filePath := volumeMountDir + "/disk0/disk-image." + diskExtension
		_, err := os.Create(filePath)

		vm, err = MapRegistryDisks(vm)
		Expect(err).ToNot(HaveOccurred())

		// verify file gets renamed by virt-handler to prevent container from
		// removing it before VM is exited
		exists, err := diskutils.FileExists(filePath)
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(false))

		// verify file rename takes place
		exists, err = diskutils.FileExists(filePath + ".virt")
		Expect(err).ToNot(HaveOccurred())
		Expect(exists).To(Equal(true))

		Expect(vm.Spec.Domain.Devices.Disks[0].Type).To(Equal("file"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Target.Device).To(Equal("vda"))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Driver.Type).To(Equal(diskExtension))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source).ToNot(Equal(nil))
		Expect(vm.Spec.Domain.Devices.Disks[0].Source.File).To(Equal(filePath + ".virt"))

		err = CleanupEphemeralDisks(vm)
		exists, err = diskutils.FileExists(volumeMountDir)
		Expect(err).ToNot(HaveOccurred())

		Expect(exists).To(Equal(false))
	}

	BeforeSuite(func() {
		err := SetLocalDirectory(tmpDir)
		if err != nil {
			panic(err)
		}
		SetLocalDataOwner(owner.Username)
	})

	AfterSuite(func() {
		os.RemoveAll(tmpDir)
	})

	Describe("registry-disk", func() {
		Context("verify helper functions", func() {
			table.DescribeTable("by verifying mapping of ",
				func(diskType string) {
					VerifyDiskType(diskType)
				},
				table.Entry("qcow2 disk", "qcow2"),
				table.Entry("raw disk", "raw"),
			)
			It("by verifying error when no disk is present", func() {

				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				vm, err := MapRegistryDisks(vm)
				Expect(err).To(HaveOccurred())
			})
			It("by verifying container generation", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vdb",
					},
				})

				containers, volumes, err := GenerateContainers(vm)
				Expect(err).ToNot(HaveOccurred())

				Expect(len(containers)).To(Equal(2))
				Expect(len(volumes)).To(Equal(2))
			})

			It("by verifying data cleanup", func() {
				vm := v1.NewMinimalVM("fake-vm")
				vm.Spec.Domain.Devices.Disks = append(vm.Spec.Domain.Devices.Disks, v1.Disk{
					Type:   "RegistryDisk:v1alpha",
					Device: "disk",
					Source: v1.DiskSource{
						Name: "someimage:v1.2.3.4",
					},
					Target: v1.DiskTarget{
						Device: "vda",
					},
				})

				volumeMountDir := generateVMBaseDir(vm)
				err = os.MkdirAll(volumeMountDir, 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk0", 0755)
				Expect(err).ToNot(HaveOccurred())
				err = os.MkdirAll(volumeMountDir+"/disk1", 0755)
				Expect(err).ToNot(HaveOccurred())

				exists, err := diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(true))

				err = CleanupEphemeralDisks(vm)
				exists, err = diskutils.FileExists(volumeMountDir)
				Expect(err).ToNot(HaveOccurred())

				Expect(exists).To(Equal(false))

			})
		})
	})
})
<|endoftext|>


<commit_before>package atlas

import (
	"encoding/json"
	"fmt"
	"github.com/sendgrid/rest"
	"log"
	"reflect"
	"strconv"
)

// MeasurementResp contains all the results of the measurements
type MeasurementResp struct {
	Measurements []int
}

// NewMeasurement create a new MeasurementRequest and fills some fields
func NewMeasurement(t string, fields map[string]string) (req *MeasurementRequest) {
	var defs []Definition

	def := NewDefinition(t, fields)
	probes := ProbeSet{
		{
			Requested: 10,
			Type:      "area",
			Value:     "WW",
			Tags:      nil,
		},
	}
	defs = append(defs, *def)
	req = &MeasurementRequest{
		Definitions: defs,
		IsOneoff:    true,
		Probes:      probes,
	}
	return
}

// NewDefinition create a new MeasuremenrRequest and fills some fields
func NewDefinition(t string, fields map[string]string) (def *Definition) {
	def = &Definition{
		Type: t,
	}
	sdef := reflect.ValueOf(&def).Elem()
	typeOfDef := sdef.Type()
	for k, v := range fields {
		// Check the field is present
		if f, ok := typeOfDef.FieldByName(k); ok {
			// Use the right type
			switch f.Name {
			case "float":
				vf, _ := strconv.ParseFloat(v, 32)
				sdef.FieldByName(k).SetFloat(vf)
			case "int":
				vi, _ := strconv.ParseInt(v, 10, 32)
				sdef.FieldByName(k).SetInt(vi)
			case "string":
				sdef.FieldByName(k).SetString(v)
			}
		}
	}
	return
}

// createMeasurement creates a measurement for all types
func createMeasurement(t string, d MeasurementRequest) (m *MeasurementResp, err error) {
	req := prepareRequest(fmt.Sprintf("measurements/%s", t))

	body, err := json.Marshal(d)
	if err != nil {
		return
	}

	req.Method = rest.Post
	req.Body = body

	log.Printf("body: %s", body)
	resp, err := rest.API(req)
	err = handleAPIResponse(resp)
	if err != nil {
		return
	}

	m = &MeasurementResp{}
	err = json.Unmarshal([]byte(resp.Body), m)
	//r, err := api.Res(base, &resp).Post(d)
	fmt.Printf("m: %v\nresp: %#v\nd: %v\n", m, string(resp.Body), d)
	if err != nil {
		err = fmt.Errorf("err: %v - m:%v", err, m)
		return
	}

	return
}

// DNS creates a measurement
func DNS(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("dns", d)
}

// HTTP creates a measurement
func HTTP(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("http", d)
}

// NTP creates a measurement
func NTP(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("ntp", d)
}

// Ping creates a measurement
func Ping(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("ping", d)
}

// SSLCert creates a measurement
func SSLCert(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("sslcert", d)
}

// Traceroute creates a measurement
func Traceroute(d MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("traceroute", d)
}
<commit_msg>Change the helpers functions for creating measurements.<commit_after>package atlas

import (
	"encoding/json"
	"fmt"
	"github.com/sendgrid/rest"
	"log"
	"reflect"
	"strconv"
)

// MeasurementResp contains all the results of the measurements
type MeasurementResp struct {
	Measurements []int
}

var (
	// If nothing is specified, use this
	defProbeSet = ProbeSet{
		{
			Requested: 10,
			Type:      "area",
			Value:     "WW",
			Tags:      nil,
		},
	}
)

// NewMeasurement create a new MeasurementRequest and fills some fields
func NewMeasurement() (req *MeasurementRequest) {
	var defs []Definition

	req = &MeasurementRequest{
		Definitions: defs,
		IsOneoff:    true,
		Probes:      defProbeSet,
	}
	return
}

func NewProbeSet(howmany int) (ps *ProbeSet) {
	ps = &ProbeSet{
		{
			Requested: howmany,
			Type:      "area",
			Value:     "WW",
			Tags:      nil,
		},
	}
	return
}

// SetParams set a few parameters in a definition list
func (d *Definition) setParams(fields map[string]string) {
	sdef := reflect.ValueOf(d).Elem()
	typeOfDef := sdef.Type()
	for k, v := range fields {
		// Check the field is present
		if f, ok := typeOfDef.FieldByName(k); ok {
			// Use the right type
			switch f.Type.Name() {
			case "float":
				vf, _ := strconv.ParseFloat(v, 32)
				sdef.FieldByName(k).SetFloat(vf)
			case "int":
				vi, _ := strconv.ParseInt(v, 10, 32)
				sdef.FieldByName(k).SetInt(vi)
			case "string":
				sdef.FieldByName(k).SetString(v)
			default:
				log.Printf("Unsupported type: %s", f.Type.Name())
			}
		}
	}
}

// AddDefinition create a new MeasurementRequest and fills some fields
func (m *MeasurementRequest) AddDefinition(fields map[string]string) (*MeasurementRequest) {
	def := new(Definition)
	def.setParams(fields)
	m.Definitions = append(m.Definitions, *def)

	return m
}

// createMeasurement creates a measurement for all types
func createMeasurement(t string, d *MeasurementRequest) (m *MeasurementResp, err error) {
	req := prepareRequest(fmt.Sprintf("measurements/%s", t))

	body, err := json.Marshal(d)
	if err != nil {
		return
	}

	req.Method = rest.Post
	req.Body = body

	log.Printf("body: %s", body)
	resp, err := rest.API(req)
	err = handleAPIResponse(resp)
	if err != nil {
		return
	}

	m = &MeasurementResp{}
	err = json.Unmarshal([]byte(resp.Body), m)
	//r, err := api.Res(base, &resp).Post(d)
	fmt.Printf("m: %v\nresp: %#v\nd: %v\n", m, string(resp.Body), d)
	if err != nil {
		err = fmt.Errorf("err: %v - m:%v", err, m)
		return
	}

	return
}

// DNS creates a measurement
func DNS(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("dns", d)
}

// HTTP creates a measurement
func HTTP(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("http", d)
}

// NTP creates a measurement
func NTP(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("ntp", d)
}

// Ping creates a measurement
func Ping(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("ping", d)
}

// SSLCert creates a measurement
func SSLCert(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("sslcert", d)
}

// Traceroute creates a measurement
func Traceroute(d *MeasurementRequest) (m *MeasurementResp, err error) {
	return createMeasurement("traceroute", d)
}
<|endoftext|>


<commit_before>package ghwebhooks

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

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	mod "github.com/influxdata/support-tools/ghWebhooks/models"
	"github.com/influxdata/telegraf/plugins/inputs"
)

func init() {
	log.Info("Starting ghWebhook server...")
	logFile, err := os.Create("server.log")
	if err != nil {
		log.WithFields(log.Fields{
			"time":  time.Now(),
			"error": err,
		}).Warn("Error in creating log file")
	}
	log.SetLevel(log.InfoLevel)
	log.SetOutput(logFile)

	inputs.Add("ghwebhooks", func() inputs.Input { return &GHWebhooks{} })
}

type GHWebhooks struct {
	ServiceAddress  string
	MeasurementName string

	sync.Mutex

	// Channel for all incoming events from github
	in   chan mod.Event
	done chan struct{}
}

func (gh *GHWebhooks) SampleConfig() string {
	return `
  # Address and port to host Webhook listener on
  service_address = ":1618"
	# Measurement name
	measurement_name = "ghWebhooks"
`
}

func (gh *GHWebhooks) Description() string {
	return "Github Webhook Event collector"
}

// Writes the points from <-gh.in to the Accumulator
func (gh *GHWebhooks) Gather(acc inputs.Accumulator) error {
	gh.Lock()
	defer gh.Unlock()
	for {
		select {
		case <-gh.done:
			return nil
		case e := <-gh.in:
			p := e.NewPoint()
			acc.Add(gh.MeasurementName, p.Fields(), p.Tags(), p.Time())
		}
	}
	return nil
}

func (gh *GHWebhooks) Start() error {
	gh.Lock()
	defer gh.Unlock()
	for {
		select {
		case <-gh.done:
			return nil
		default:
			r := mux.NewRouter()
			r.HandleFunc("/webhooks", gh.webhookHandler).Methods("POST")
			http.ListenAndServe(fmt.Sprintf(":%s", gh.ServiceAddress), r)
		}
	}
}

func (gh *GHWebhooks) Stop() {
	gh.Lock()
	defer gh.Unlock()
	log.Println("Stopping the ghWebhooks service")
	close(gh.done)
	close(gh.in)
}

// Handles the /webhooks route
func (gh *GHWebhooks) webhookHandler(w http.ResponseWriter, r *http.Request) {
	eventType := r.Header["X-Github-Event"][0]
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": eventType, "error": err}
		log.WithFields(fields).Fatal("Error reading Github payload")
	}

	// Send event down chan to GHWebhooks
	e := NewEvent(data, eventType)
	gh.in <- e
	fmt.Printf("%v\n", e.NewPoint())
	w.WriteHeader(http.StatusOK)
}

func newCommitComment(data []byte) mod.Event {
	commitCommentStruct := mod.CommitCommentEvent{}
	err := json.Unmarshal(data, &commitCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "CommitCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return commitCommentStruct
}

func newCreate(data []byte) mod.Event {
	createStruct := mod.CreateEvent{}
	err := json.Unmarshal(data, &createStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "CreateEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return createStruct
}

func newDelete(data []byte) mod.Event {
	deleteStruct := mod.DeleteEvent{}
	err := json.Unmarshal(data, &deleteStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeleteEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deleteStruct
}

func newDeployment(data []byte) mod.Event {
	deploymentStruct := mod.DeploymentEvent{}
	err := json.Unmarshal(data, &deploymentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeploymentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deploymentStruct
}

func newDeploymentStatus(data []byte) mod.Event {
	deploymentStatusStruct := mod.DeploymentStatusEvent{}
	err := json.Unmarshal(data, &deploymentStatusStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeploymentStatusEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deploymentStatusStruct
}

func newFork(data []byte) mod.Event {
	forkStruct := mod.ForkEvent{}
	err := json.Unmarshal(data, &forkStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "ForkEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return forkStruct
}

func newGollum(data []byte) mod.Event {
	gollumStruct := mod.GollumEvent{}
	err := json.Unmarshal(data, &gollumStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "GollumEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return gollumStruct
}

func newIssueComment(data []byte) mod.Event {
	issueCommentStruct := mod.IssueCommentEvent{}
	err := json.Unmarshal(data, &issueCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "IssueCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return issueCommentStruct
}

func newIssues(data []byte) mod.Event {
	issuesStruct := mod.IssuesEvent{}
	err := json.Unmarshal(data, &issuesStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "IssuesEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return issuesStruct
}

func newMember(data []byte) mod.Event {
	memberStruct := mod.MemberEvent{}
	err := json.Unmarshal(data, &memberStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "MemberEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return memberStruct
}

func newMembership(data []byte) mod.Event {
	membershipStruct := mod.MembershipEvent{}
	err := json.Unmarshal(data, &membershipStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "MembershipEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return membershipStruct
}

func newPageBuild(data []byte) mod.Event {
	pageBuildEvent := mod.PageBuildEvent{}
	err := json.Unmarshal(data, &pageBuildEvent)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PageBuildEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pageBuildEvent
}

func newPublic(data []byte) mod.Event {
	publicEvent := mod.PublicEvent{}
	err := json.Unmarshal(data, &publicEvent)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PublicEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return publicEvent
}

func newPullRequest(data []byte) mod.Event {
	pullRequestStruct := mod.PullRequestEvent{}
	err := json.Unmarshal(data, &pullRequestStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PullRequestEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pullRequestStruct
}

func newPullRequestReviewComment(data []byte) mod.Event {
	pullRequestReviewCommentStruct := mod.PullRequestReviewCommentEvent{}
	err := json.Unmarshal(data, &pullRequestReviewCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PullRequestReviewCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pullRequestReviewCommentStruct
}

func newPush(data []byte) mod.Event {
	pushStruct := mod.PushEvent{}
	err := json.Unmarshal(data, &pushStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PushEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pushStruct
}

func newRelease(data []byte) mod.Event {
	releaseStruct := mod.ReleaseEvent{}
	err := json.Unmarshal(data, &releaseStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "ReleaseEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return releaseStruct
}

func newRepository(data []byte) mod.Event {
	repositoryStruct := mod.RepositoryEvent{}
	err := json.Unmarshal(data, &repositoryStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "RepositoryEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return repositoryStruct
}

func newStatus(data []byte) mod.Event {
	statusStruct := mod.StatusEvent{}
	err := json.Unmarshal(data, &statusStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "StatusEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return statusStruct
}

func newTeamAdd(data []byte) mod.Event {
	teamAddStruct := mod.TeamAddEvent{}
	err := json.Unmarshal(data, &teamAddStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "TeamAddEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return teamAddStruct
}

func newWatch(data []byte) mod.Event {
	watchStruct := mod.WatchEvent{}
	err := json.Unmarshal(data, &watchStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "WatchEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return watchStruct
}

func NewEvent(r []byte, t string) mod.Event {
	log.WithFields(log.Fields{"event": t, "time": time.Now()}).Info("Event Recieved")
	switch t {
	case "commit_comment":
		return newCommitComment(r)
	case "create":
		return newCreate(r)
	case "delete":
		return newDelete(r)
	case "deployment":
		return newDeployment(r)
	case "deployment_status":
		return newDeploymentStatus(r)
	case "fork":
		return newFork(r)
	case "gollum":
		return newGollum(r)
	case "issue_comment":
		return newIssueComment(r)
	case "issues":
		return newIssues(r)
	case "member":
		return newMember(r)
	case "membership":
		return newMembership(r)
	case "page_build":
		return newPageBuild(r)
	case "public":
		return newPublic(r)
	case "pull_request":
		return newPullRequest(r)
	case "pull_request_review_comment":
		return newPullRequestReviewComment(r)
	case "push":
		return newPush(r)
	case "release":
		return newRelease(r)
	case "repository":
		return newRepository(r)
	case "status":
		return newStatus(r)
	case "team_add":
		return newTeamAdd(r)
	case "watch":
		return newWatch(r)
	}
	return nil
}
<commit_msg>Change start implementation<commit_after>package ghwebhooks

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

	log "github.com/Sirupsen/logrus"
	"github.com/gorilla/mux"
	mod "github.com/influxdata/support-tools/ghWebhooks/models"
	"github.com/influxdata/telegraf/plugins/inputs"
)

func init() {
	log.Info("Starting ghWebhook server...")
	logFile, err := os.Create("server.log")
	if err != nil {
		log.WithFields(log.Fields{
			"time":  time.Now(),
			"error": err,
		}).Warn("Error in creating log file")
	}
	log.SetLevel(log.InfoLevel)
	log.SetOutput(logFile)

	inputs.Add("ghwebhooks", func() inputs.Input { return &GHWebhooks{} })
}

type GHWebhooks struct {
	ServiceAddress  string
	MeasurementName string

	sync.Mutex

	// Channel for all incoming events from github
	in   chan mod.Event
	done chan struct{}
}

func (gh *GHWebhooks) SampleConfig() string {
	return `
  # Address and port to host Webhook listener on
  service_address = ":1618"
	# Measurement name
	measurement_name = "ghWebhooks"
`
}

func (gh *GHWebhooks) Description() string {
	return "Github Webhook Event collector"
}

// Writes the points from <-gh.in to the Accumulator
func (gh *GHWebhooks) Gather(acc inputs.Accumulator) error {
	gh.Lock()
	defer gh.Unlock()
	for {
		select {
		case <-gh.done:
			return nil
		case e := <-gh.in:
			p := e.NewPoint()
			acc.Add(gh.MeasurementName, p.Fields(), p.Tags(), p.Time())
		}
	}
	return nil
}

func (gh *GHWebhooks) listen() error {
	r := mux.NewRouter()
	r.HandleFunc("/webhooks", gh.webhookHandler).Methods("POST")
	http.ListenAndServe(fmt.Sprintf(":%s", gh.ServiceAddress), r)
}

func (gh *GHWebhooks) Start() error {
	gh.done = make(chan struct{})
	gh.in = make(chan mod.Event)
	// Start the UDP listener
	go gh.listen()
	// Start the line parser
	log.Printf("Started the ghwebhooks service on %s\n", s.ServiceAddress)
}

func (gh *GHWebhooks) Stop() {
	gh.Lock()
	defer gh.Unlock()
	log.Println("Stopping the ghWebhooks service")
	close(gh.done)
	close(gh.in)
}

// Handles the /webhooks route
func (gh *GHWebhooks) webhookHandler(w http.ResponseWriter, r *http.Request) {
	gh.Lock()
	defer gh.Unlock()
	eventType := r.Header["X-Github-Event"][0]
	data, err := ioutil.ReadAll(r.Body)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": eventType, "error": err}
		log.WithFields(fields).Fatal("Error reading Github payload")
	}

	// Send event down chan to GHWebhooks
	e := NewEvent(data, eventType)
	gh.in <- e
	fmt.Printf("%v\n", e.NewPoint())
	w.WriteHeader(http.StatusOK)
}

func newCommitComment(data []byte) mod.Event {
	commitCommentStruct := mod.CommitCommentEvent{}
	err := json.Unmarshal(data, &commitCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "CommitCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return commitCommentStruct
}

func newCreate(data []byte) mod.Event {
	createStruct := mod.CreateEvent{}
	err := json.Unmarshal(data, &createStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "CreateEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return createStruct
}

func newDelete(data []byte) mod.Event {
	deleteStruct := mod.DeleteEvent{}
	err := json.Unmarshal(data, &deleteStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeleteEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deleteStruct
}

func newDeployment(data []byte) mod.Event {
	deploymentStruct := mod.DeploymentEvent{}
	err := json.Unmarshal(data, &deploymentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeploymentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deploymentStruct
}

func newDeploymentStatus(data []byte) mod.Event {
	deploymentStatusStruct := mod.DeploymentStatusEvent{}
	err := json.Unmarshal(data, &deploymentStatusStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "DeploymentStatusEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return deploymentStatusStruct
}

func newFork(data []byte) mod.Event {
	forkStruct := mod.ForkEvent{}
	err := json.Unmarshal(data, &forkStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "ForkEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return forkStruct
}

func newGollum(data []byte) mod.Event {
	gollumStruct := mod.GollumEvent{}
	err := json.Unmarshal(data, &gollumStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "GollumEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return gollumStruct
}

func newIssueComment(data []byte) mod.Event {
	issueCommentStruct := mod.IssueCommentEvent{}
	err := json.Unmarshal(data, &issueCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "IssueCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return issueCommentStruct
}

func newIssues(data []byte) mod.Event {
	issuesStruct := mod.IssuesEvent{}
	err := json.Unmarshal(data, &issuesStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "IssuesEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return issuesStruct
}

func newMember(data []byte) mod.Event {
	memberStruct := mod.MemberEvent{}
	err := json.Unmarshal(data, &memberStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "MemberEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return memberStruct
}

func newMembership(data []byte) mod.Event {
	membershipStruct := mod.MembershipEvent{}
	err := json.Unmarshal(data, &membershipStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "MembershipEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return membershipStruct
}

func newPageBuild(data []byte) mod.Event {
	pageBuildEvent := mod.PageBuildEvent{}
	err := json.Unmarshal(data, &pageBuildEvent)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PageBuildEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pageBuildEvent
}

func newPublic(data []byte) mod.Event {
	publicEvent := mod.PublicEvent{}
	err := json.Unmarshal(data, &publicEvent)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PublicEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return publicEvent
}

func newPullRequest(data []byte) mod.Event {
	pullRequestStruct := mod.PullRequestEvent{}
	err := json.Unmarshal(data, &pullRequestStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PullRequestEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pullRequestStruct
}

func newPullRequestReviewComment(data []byte) mod.Event {
	pullRequestReviewCommentStruct := mod.PullRequestReviewCommentEvent{}
	err := json.Unmarshal(data, &pullRequestReviewCommentStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PullRequestReviewCommentEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pullRequestReviewCommentStruct
}

func newPush(data []byte) mod.Event {
	pushStruct := mod.PushEvent{}
	err := json.Unmarshal(data, &pushStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "PushEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return pushStruct
}

func newRelease(data []byte) mod.Event {
	releaseStruct := mod.ReleaseEvent{}
	err := json.Unmarshal(data, &releaseStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "ReleaseEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return releaseStruct
}

func newRepository(data []byte) mod.Event {
	repositoryStruct := mod.RepositoryEvent{}
	err := json.Unmarshal(data, &repositoryStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "RepositoryEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return repositoryStruct
}

func newStatus(data []byte) mod.Event {
	statusStruct := mod.StatusEvent{}
	err := json.Unmarshal(data, &statusStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "StatusEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return statusStruct
}

func newTeamAdd(data []byte) mod.Event {
	teamAddStruct := mod.TeamAddEvent{}
	err := json.Unmarshal(data, &teamAddStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "TeamAddEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return teamAddStruct
}

func newWatch(data []byte) mod.Event {
	watchStruct := mod.WatchEvent{}
	err := json.Unmarshal(data, &watchStruct)
	if err != nil {
		fields := log.Fields{"time": time.Now(), "event": "WatchEvent", "error": err}
		log.WithFields(fields).Fatalf("Error in unmarshaling JSON")
	}
	return watchStruct
}

func NewEvent(r []byte, t string) mod.Event {
	log.WithFields(log.Fields{"event": t, "time": time.Now()}).Info("Event Recieved")
	switch t {
	case "commit_comment":
		return newCommitComment(r)
	case "create":
		return newCreate(r)
	case "delete":
		return newDelete(r)
	case "deployment":
		return newDeployment(r)
	case "deployment_status":
		return newDeploymentStatus(r)
	case "fork":
		return newFork(r)
	case "gollum":
		return newGollum(r)
	case "issue_comment":
		return newIssueComment(r)
	case "issues":
		return newIssues(r)
	case "member":
		return newMember(r)
	case "membership":
		return newMembership(r)
	case "page_build":
		return newPageBuild(r)
	case "public":
		return newPublic(r)
	case "pull_request":
		return newPullRequest(r)
	case "pull_request_review_comment":
		return newPullRequestReviewComment(r)
	case "push":
		return newPush(r)
	case "release":
		return newRelease(r)
	case "repository":
		return newRepository(r)
	case "status":
		return newStatus(r)
	case "team_add":
		return newTeamAdd(r)
	case "watch":
		return newWatch(r)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Cockroach 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.
//
// Author: Tobias Schottdorf (tobias.schottdorf@gmail.com)

package acceptance

import (
	"fmt"
	"math/rand"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/acceptance/cluster"
	"github.com/cockroachdb/cockroach/pkg/util/caller"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/pkg/errors"
)

func TestPartitionNemesis(t *testing.T) {
	t.Skip("only enabled for manually playing with the partitioning agent")
	SkipUnlessLocal(t)
	runTestOnConfigs(t, func(ctx context.Context, t *testing.T, c cluster.Cluster, cfg cluster.TestConfig) {
		stopper.RunWorker(func() {
			BidirectionalPartitionNemesis(ctx, t, c, stopper)
		})
		select {
		case <-time.After(*flagDuration):
		case <-stopper.ShouldStop():
		}
	})
}

func TestPartitionBank(t *testing.T) {
	t.Skip("#7978")
	SkipUnlessPrivileged(t)
	runTestOnConfigs(t, testBankWithNemesis(BidirectionalPartitionNemesis))
}

type Bank struct {
	cluster.Cluster
	*testing.T
	accounts, initialBalance int
}

func (b *Bank) must(err error) {
	if err != nil {
		f, l, _ := caller.Lookup(1)
		b.Fatal(errors.Wrapf(err, "%s:%d", f, l))
	}
}

// NewBank creates a Bank.
// TODO(tamird,tschottdorf): share this code with other bank test(s).
func NewBank(t *testing.T, c cluster.Cluster) *Bank {
	return &Bank{Cluster: c, T: t}
}

func (b *Bank) exec(ctx context.Context, query string, vars ...interface{}) error {
	db := makePGClient(b.T, b.PGUrl(ctx, 0))
	defer db.Close()
	_, err := db.Exec(query, vars...)
	return err
}

// Init sets up the bank for the given number of accounts, each of which
// receiving a deposit of the given amount.
// This should be called before any nemeses are active; it will fail the test
// if unsuccessful.
func (b *Bank) Init(ctx context.Context, numAccounts, initialBalance int) {
	b.accounts = numAccounts
	b.initialBalance = initialBalance

	b.must(b.exec(ctx, `CREATE DATABASE IF NOT EXISTS bank`))
	b.must(b.exec(ctx, `DROP TABLE IF EXISTS bank.accounts`))
	const schema = `CREATE TABLE bank.accounts (id INT PRIMARY KEY, balance INT NOT NULL)`
	b.must(b.exec(ctx, schema))
	for i := 0; i < numAccounts; i++ {
		b.must(b.exec(ctx, `INSERT INTO bank.accounts (id, balance) VALUES ($1, $2)`,
			i, initialBalance))
	}
}

// Verify makes sure that the total amount of money in the system has not
// changed.
func (b *Bank) Verify(ctx context.Context) {
	exp := b.accounts * b.initialBalance
	db := makePGClient(b.T, b.PGUrl(ctx, 0))
	defer db.Close()
	r := db.QueryRow(`SELECT SUM(balance) FROM bank.accounts`)
	var act int
	b.must(r.Scan(&act))
	if act != exp {
		b.Fatalf("bank is worth $%d, should be $%d", act, exp)
	}
}

func (b *Bank) logFailed(i int, v interface{}) {
	log.Warningf(context.Background(), "%d: %v", i, v)
}
func (b *Bank) logBegin(i int, from, to, amount int) {
	log.Warningf(context.Background(), "%d: %d trying to give $%d to %d", i, from, amount, to)
}
func (b *Bank) logSuccess(i int, from, to, amount int) {
	log.Warningf(context.Background(), "%d: %d gave $%d to %d", i, from, amount, to)
}

// Invoke transfers a random amount of money between random accounts.
func (b *Bank) Invoke(ctx context.Context, i int) {
	handle := func(err error) {
		if err != nil {
			panic(err)
		}
	}
	var from, to int
	{
		p := rand.Perm(b.accounts)
		from, to = p[0], p[1]
	}
	amount := rand.Intn(b.initialBalance)
	b.logBegin(i, from, to, amount)
	defer func() {
		if r := recover(); r != nil {
			b.logFailed(i, r)
		} else {
			b.logSuccess(i, from, to, amount)
		}
	}()

	db := makePGClient(b.T, b.PGUrl(ctx, i%b.NumNodes()))
	defer db.Close()
	txn, err := db.Begin()
	handle(err)
	// The following SQL queries are intentionally unoptimized.
	var bFrom, bTo int
	{
		rFrom := txn.QueryRow(`SELECT balance FROM bank.accounts WHERE id = $1`, from)
		handle(rFrom.Scan(&bFrom))
		rTo := txn.QueryRow(`SELECT balance FROM bank.accounts WHERE id = $1`, to)
		handle(rTo.Scan(&bTo))
	}
	if diff := bFrom - amount; diff < 0 {
		handle(fmt.Errorf("%d is %d short to pay $%d", bFrom, -diff, amount))
	}
	_, err = txn.Exec(`UPDATE bank.accounts SET balance = $1 WHERE id = $2`, bFrom-amount, from)
	handle(err)
	_, err = txn.Exec(`UPDATE bank.accounts SET balance = $1 WHERE id = $2`, bTo+amount, to)
	handle(err)
}

func testBankWithNemesis(nemeses ...NemesisFn) configTestRunner {
	return func(ctx context.Context, t *testing.T, c cluster.Cluster, cfg cluster.TestConfig) {
		const (
			concurrency = 5
			accounts    = 10
		)
		deadline := timeutil.Now().Add(cfg.Duration)
		b := NewBank(t, c)
		b.Init(ctx, accounts, 10)
		for _, nemesis := range nemeses {
			stopper.RunWorker(func() {
				nemesis(ctx, t, c, stopper)
			})
		}
		for i := 0; i < concurrency; i++ {
			localI := i
			if err := stopper.RunAsyncTask(ctx, func(_ context.Context) {
				for timeutil.Now().Before(deadline) {
					select {
					case <-stopper.ShouldQuiesce():
						return
					default:
					}
					b.Invoke(ctx, localI)
				}
			}); err != nil {
				t.Fatal(err)
			}
		}
		select {
		case <-stopper.ShouldStop():
		case <-time.After(cfg.Duration):
		}
		log.Warningf(ctx, "finishing test")
		b.Verify(ctx)
	}
}
<commit_msg>acceptance: fix and re-enable TestPartitionBank<commit_after>// Copyright 2016 The Cockroach 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.
//
// Author: Tobias Schottdorf (tobias.schottdorf@gmail.com)

package acceptance

import (
	"fmt"
	"math/rand"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/cockroachdb/cockroach/pkg/acceptance/cluster"
	"github.com/cockroachdb/cockroach/pkg/util/caller"
	"github.com/cockroachdb/cockroach/pkg/util/log"
	"github.com/cockroachdb/cockroach/pkg/util/stop"
	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
	"github.com/pkg/errors"
)

func TestPartitionNemesis(t *testing.T) {
	t.Skip("only enabled for manually playing with the partitioning agent")
	SkipUnlessLocal(t)
	runTestOnConfigs(t, func(ctx context.Context, t *testing.T, c cluster.Cluster, cfg cluster.TestConfig) {
		stopper.RunWorker(func() {
			BidirectionalPartitionNemesis(ctx, t, c, stopper)
		})
		select {
		case <-time.After(*flagDuration):
		case <-stopper.ShouldStop():
		}
	})
}

func TestPartitionBank(t *testing.T) {
	SkipUnlessPrivileged(t)
	runTestOnConfigs(t, testBankWithNemesis(BidirectionalPartitionNemesis))
}

type Bank struct {
	cluster.Cluster
	*testing.T
	accounts, initialBalance int
}

func (b *Bank) must(err error) {
	if err != nil {
		f, l, _ := caller.Lookup(1)
		b.Fatal(errors.Wrapf(err, "%s:%d", f, l))
	}
}

// NewBank creates a Bank.
// TODO(tamird,tschottdorf): share this code with other bank test(s).
func NewBank(t *testing.T, c cluster.Cluster) *Bank {
	return &Bank{Cluster: c, T: t}
}

func (b *Bank) exec(ctx context.Context, query string, vars ...interface{}) error {
	db := makePGClient(b.T, b.PGUrl(ctx, 0))
	defer db.Close()
	_, err := db.Exec(query, vars...)
	return err
}

// Init sets up the bank for the given number of accounts, each of which
// receiving a deposit of the given amount.
// This should be called before any nemeses are active; it will fail the test
// if unsuccessful.
func (b *Bank) Init(ctx context.Context, numAccounts, initialBalance int) {
	b.accounts = numAccounts
	b.initialBalance = initialBalance

	b.must(b.exec(ctx, `CREATE DATABASE IF NOT EXISTS bank`))
	b.must(b.exec(ctx, `DROP TABLE IF EXISTS bank.accounts`))
	const schema = `CREATE TABLE bank.accounts (id INT PRIMARY KEY, balance INT NOT NULL)`
	b.must(b.exec(ctx, schema))
	for i := 0; i < numAccounts; i++ {
		b.must(b.exec(ctx, `INSERT INTO bank.accounts (id, balance) VALUES ($1, $2)`,
			i, initialBalance))
	}
}

// Verify makes sure that the total amount of money in the system has not
// changed.
func (b *Bank) Verify(ctx context.Context) {
	log.Info(ctx, "verifying")
	exp := b.accounts * b.initialBalance
	db := makePGClient(b.T, b.PGUrl(ctx, 0))
	defer db.Close()
	r := db.QueryRow(`SELECT SUM(balance) FROM bank.accounts`)
	var act int
	b.must(r.Scan(&act))
	if act != exp {
		b.Fatalf("bank is worth $%d, should be $%d", act, exp)
	}
}

func (b *Bank) logFailed(i int, v interface{}) {
	log.Warningf(context.Background(), "%d: %v", i, v)
}
func (b *Bank) logBegin(i int, from, to, amount int) {
	log.Warningf(context.Background(), "%d: %d trying to give $%d to %d", i, from, amount, to)
}
func (b *Bank) logSuccess(i int, from, to, amount int) {
	log.Warningf(context.Background(), "%d: %d gave $%d to %d", i, from, amount, to)
}

// Invoke transfers a random amount of money between random accounts.
func (b *Bank) Invoke(ctx context.Context, i int) {
	handle := func(err error) {
		if err != nil {
			panic(err)
		}
	}
	var from, to int
	{
		p := rand.Perm(b.accounts)
		from, to = p[0], p[1]
	}
	amount := rand.Intn(b.initialBalance)
	b.logBegin(i, from, to, amount)
	defer func() {
		if r := recover(); r != nil {
			b.logFailed(i, r)
		} else {
			b.logSuccess(i, from, to, amount)
		}
	}()

	db := makePGClient(b.T, b.PGUrl(ctx, i%b.NumNodes()))
	defer db.Close()
	txn, err := db.Begin()
	handle(err)
	// The following SQL queries are intentionally unoptimized.
	var bFrom, bTo int
	{
		rFrom := txn.QueryRow(`SELECT balance FROM bank.accounts WHERE id = $1`, from)
		handle(rFrom.Scan(&bFrom))
		rTo := txn.QueryRow(`SELECT balance FROM bank.accounts WHERE id = $1`, to)
		handle(rTo.Scan(&bTo))
	}
	if diff := bFrom - amount; diff < 0 {
		handle(fmt.Errorf("%d is %d short to pay $%d", bFrom, -diff, amount))
	}
	_, err = txn.Exec(`UPDATE bank.accounts SET balance = $1 WHERE id = $2`, bFrom-amount, from)
	handle(err)
	_, err = txn.Exec(`UPDATE bank.accounts SET balance = $1 WHERE id = $2`, bTo+amount, to)
	handle(err)
	handle(txn.Commit())
}

func testBankWithNemesis(nemeses ...NemesisFn) configTestRunner {
	return func(ctx context.Context, t *testing.T, c cluster.Cluster, cfg cluster.TestConfig) {
		const accounts = 10
		b := NewBank(t, c)
		b.Init(ctx, accounts, 10)
		runTransactionsAndNemeses(ctx, t, c, b, cfg.Duration, nemeses...)
		b.Verify(ctx)
	}
}

func runTransactionsAndNemeses(
	ctx context.Context,
	t *testing.T,
	c cluster.Cluster,
	b *Bank,
	duration time.Duration,
	nemeses ...NemesisFn,
) {
	deadline := timeutil.Now().Add(duration)
	// We're going to run the nemeses for the duration of this function, which may
	// return before `stopper` is stopped.
	nemesesStopper := stop.NewStopper()
	defer nemesesStopper.Stop()
	const concurrency = 5
	for _, nemesis := range nemeses {
		stopper.RunWorker(func() {
			nemesis(ctx, t, c, nemesesStopper)
		})
	}
	for i := 0; i < concurrency; i++ {
		localI := i
		if err := stopper.RunAsyncTask(ctx, func(_ context.Context) {
			for timeutil.Now().Before(deadline) {
				select {
				case <-stopper.ShouldQuiesce():
					return
				default:
				}
				b.Invoke(ctx, localI)
			}
		}); err != nil {
			t.Fatal(err)
		}
	}
	select {
	case <-stopper.ShouldStop():
	case <-time.After(duration):
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package cache

import (
	"testing"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/extensions"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/sets"
)

func TestStoreToNodeLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	ids := sets.NewString("foo", "bar", "baz")
	for id := range ids {
		store.Add(&api.Node{ObjectMeta: api.ObjectMeta{Name: id}})
	}
	sml := StoreToNodeLister{store}

	gotNodes, err := sml.List()
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	got := make([]string, len(gotNodes.Items))
	for ix := range gotNodes.Items {
		got[ix] = gotNodes.Items[ix].Name
	}
	if !ids.HasAll(got...) || len(got) != len(ids) {
		t.Errorf("Expected %v, got %v", ids, got)
	}
}

func TestStoreToReplicationControllerLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	lister := StoreToReplicationControllerLister{store}
	testCases := []struct {
		inRCs      []*api.ReplicationController
		list       func() ([]api.ReplicationController, error)
		outRCNames sets.String
		expectErr  bool
	}{
		// Basic listing with all labels and no selectors
		{
			inRCs: []*api.ReplicationController{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
			},
			list: func() ([]api.ReplicationController, error) {
				return lister.List()
			},
			outRCNames: sets.NewString("basic"),
		},
		// No pod labels
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "baz"},
					},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{Name: "pod1", Namespace: "ns"},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString(),
			expectErr:  true,
		},
		// No RC selectors
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Namespace: "ns",
						Labels:    map[string]string{"foo": "bar"},
					},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString(),
			expectErr:  true,
		},
		// Matching labels to selectors and namespace
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "foo"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
				{
					ObjectMeta: api.ObjectMeta{Name: "bar", Namespace: "ns"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Labels:    map[string]string{"foo": "bar"},
						Namespace: "ns",
					},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString("bar"),
		},
	}
	for _, c := range testCases {
		for _, r := range c.inRCs {
			store.Add(r)
		}

		gotControllers, err := c.list()
		if err != nil && c.expectErr {
			continue
		} else if c.expectErr {
			t.Fatalf("Expected error, got none")
		} else if err != nil {
			t.Fatalf("Unexpected error %#v", err)
		}
		gotNames := make([]string, len(gotControllers))
		for ix := range gotControllers {
			gotNames[ix] = gotControllers[ix].Name
		}
		if !c.outRCNames.HasAll(gotNames...) || len(gotNames) != len(c.outRCNames) {
			t.Errorf("Unexpected got controllers %+v expected %+v", gotNames, c.outRCNames)
		}
	}
}

func TestStoreToDaemonSetLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	lister := StoreToDaemonSetLister{store}
	testCases := []struct {
		inDSs             []*extensions.DaemonSet
		list              func() ([]extensions.DaemonSet, error)
		outDaemonSetNames sets.String
		expectErr         bool
	}{
		// Basic listing
		{
			inDSs: []*extensions.DaemonSet{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
			},
			list: func() ([]extensions.DaemonSet, error) {
				return lister.List()
			},
			outDaemonSetNames: sets.NewString("basic"),
		},
		// Listing multiple daemon sets
		{
			inDSs: []*extensions.DaemonSet{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
				{ObjectMeta: api.ObjectMeta{Name: "complex"}},
				{ObjectMeta: api.ObjectMeta{Name: "complex2"}},
			},
			list: func() ([]extensions.DaemonSet, error) {
				return lister.List()
			},
			outDaemonSetNames: sets.NewString("basic", "complex", "complex2"),
		},
		// No pod labels
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "baz"},
					},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{Name: "pod1", Namespace: "ns"},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString(),
			expectErr:         true,
		},
		// No DS selectors
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Namespace: "ns",
						Labels:    map[string]string{"foo": "bar"},
					},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString(),
			expectErr:         true,
		},
		// Matching labels to selectors and namespace
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "foo"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
				{
					ObjectMeta: api.ObjectMeta{Name: "bar", Namespace: "ns"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Labels:    map[string]string{"foo": "bar"},
						Namespace: "ns",
					},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString("bar"),
		},
	}
	for _, c := range testCases {
		for _, r := range c.inDSs {
			store.Add(r)
		}

		daemonSets, err := c.list()
		if err != nil && c.expectErr {
			continue
		} else if c.expectErr {
			t.Fatalf("Expected error, got none")
		} else if err != nil {
			t.Fatalf("Unexpected error %#v", err)
		}
		daemonSetNames := make([]string, len(daemonSets))
		for ix := range daemonSets {
			daemonSetNames[ix] = daemonSets[ix].Name
		}
		if !c.outDaemonSetNames.HasAll(daemonSetNames...) || len(daemonSetNames) != len(c.outDaemonSetNames) {
			t.Errorf("Unexpected got controllers %+v expected %+v", daemonSetNames, c.outDaemonSetNames)
		}
	}
}

func TestStoreToPodLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	ids := []string{"foo", "bar", "baz"}
	for _, id := range ids {
		store.Add(&api.Pod{
			ObjectMeta: api.ObjectMeta{
				Name:   id,
				Labels: map[string]string{"name": id},
			},
		})
	}
	spl := StoreToPodLister{store}

	for _, id := range ids {
		got, err := spl.List(labels.Set{"name": id}.AsSelector())
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
			continue
		}
		if e, a := 1, len(got); e != a {
			t.Errorf("Expected %v, got %v", e, a)
			continue
		}
		if e, a := id, got[0].Name; e != a {
			t.Errorf("Expected %v, got %v", e, a)
			continue
		}

		exists, err := spl.Exists(&api.Pod{ObjectMeta: api.ObjectMeta{Name: id}})
		if err != nil {
			t.Errorf("unexpected error: %v", err)
		}
		if !exists {
			t.Errorf("exists returned false for %v", id)
		}
	}

	exists, err := spl.Exists(&api.Pod{ObjectMeta: api.ObjectMeta{Name: "qux"}})
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}
	if exists {
		t.Errorf("Unexpected pod exists")
	}
}

func TestStoreToServiceLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	store.Add(&api.Service{
		ObjectMeta: api.ObjectMeta{Name: "foo"},
		Spec: api.ServiceSpec{
			Selector: map[string]string{},
		},
	})
	store.Add(&api.Service{ObjectMeta: api.ObjectMeta{Name: "bar"}})
	ssl := StoreToServiceLister{store}

	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Name:   "foopod",
			Labels: map[string]string{"role": "foo"},
		},
	}

	services, err := ssl.GetPodServices(pod)
	if err != nil {
		t.Errorf("Unexpected error: %v", err)
	}

	if len(services) != 1 {
		t.Fatalf("Expected 1 service, got %v", len(services))
	}
	if e, a := "foo", services[0].Name; e != a {
		t.Errorf("Expected service %q, got %q", e, a)
	}
}
<commit_msg>replace Fatalf to Errorf<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

package cache

import (
	"testing"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/extensions"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/pkg/util/sets"
)

func TestStoreToNodeLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	ids := sets.NewString("foo", "bar", "baz")
	for id := range ids {
		store.Add(&api.Node{ObjectMeta: api.ObjectMeta{Name: id}})
	}
	sml := StoreToNodeLister{store}

	gotNodes, err := sml.List()
	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	got := make([]string, len(gotNodes.Items))
	for ix := range gotNodes.Items {
		got[ix] = gotNodes.Items[ix].Name
	}
	if !ids.HasAll(got...) || len(got) != len(ids) {
		t.Errorf("Expected %v, got %v", ids, got)
	}
}

func TestStoreToReplicationControllerLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	lister := StoreToReplicationControllerLister{store}
	testCases := []struct {
		inRCs      []*api.ReplicationController
		list       func() ([]api.ReplicationController, error)
		outRCNames sets.String
		expectErr  bool
	}{
		// Basic listing with all labels and no selectors
		{
			inRCs: []*api.ReplicationController{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
			},
			list: func() ([]api.ReplicationController, error) {
				return lister.List()
			},
			outRCNames: sets.NewString("basic"),
		},
		// No pod labels
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "baz"},
					},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{Name: "pod1", Namespace: "ns"},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString(),
			expectErr:  true,
		},
		// No RC selectors
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Namespace: "ns",
						Labels:    map[string]string{"foo": "bar"},
					},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString(),
			expectErr:  true,
		},
		// Matching labels to selectors and namespace
		{
			inRCs: []*api.ReplicationController{
				{
					ObjectMeta: api.ObjectMeta{Name: "foo"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
				{
					ObjectMeta: api.ObjectMeta{Name: "bar", Namespace: "ns"},
					Spec: api.ReplicationControllerSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
			},
			list: func() ([]api.ReplicationController, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Labels:    map[string]string{"foo": "bar"},
						Namespace: "ns",
					},
				}
				return lister.GetPodControllers(pod)
			},
			outRCNames: sets.NewString("bar"),
		},
	}
	for _, c := range testCases {
		for _, r := range c.inRCs {
			store.Add(r)
		}

		gotControllers, err := c.list()
		if err != nil && c.expectErr {
			continue
		} else if c.expectErr {
			t.Error("Expected error, got none")
			continue
		} else if err != nil {
			t.Errorf("Unexpected error %#v", err)
			continue
		}
		gotNames := make([]string, len(gotControllers))
		for ix := range gotControllers {
			gotNames[ix] = gotControllers[ix].Name
		}
		if !c.outRCNames.HasAll(gotNames...) || len(gotNames) != len(c.outRCNames) {
			t.Errorf("Unexpected got controllers %+v expected %+v", gotNames, c.outRCNames)
		}
	}
}

func TestStoreToDaemonSetLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	lister := StoreToDaemonSetLister{store}
	testCases := []struct {
		inDSs             []*extensions.DaemonSet
		list              func() ([]extensions.DaemonSet, error)
		outDaemonSetNames sets.String
		expectErr         bool
	}{
		// Basic listing
		{
			inDSs: []*extensions.DaemonSet{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
			},
			list: func() ([]extensions.DaemonSet, error) {
				return lister.List()
			},
			outDaemonSetNames: sets.NewString("basic"),
		},
		// Listing multiple daemon sets
		{
			inDSs: []*extensions.DaemonSet{
				{ObjectMeta: api.ObjectMeta{Name: "basic"}},
				{ObjectMeta: api.ObjectMeta{Name: "complex"}},
				{ObjectMeta: api.ObjectMeta{Name: "complex2"}},
			},
			list: func() ([]extensions.DaemonSet, error) {
				return lister.List()
			},
			outDaemonSetNames: sets.NewString("basic", "complex", "complex2"),
		},
		// No pod labels
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "baz"},
					},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{Name: "pod1", Namespace: "ns"},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString(),
			expectErr:         true,
		},
		// No DS selectors
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "basic", Namespace: "ns"},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Namespace: "ns",
						Labels:    map[string]string{"foo": "bar"},
					},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString(),
			expectErr:         true,
		},
		// Matching labels to selectors and namespace
		{
			inDSs: []*extensions.DaemonSet{
				{
					ObjectMeta: api.ObjectMeta{Name: "foo"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
				{
					ObjectMeta: api.ObjectMeta{Name: "bar", Namespace: "ns"},
					Spec: extensions.DaemonSetSpec{
						Selector: map[string]string{"foo": "bar"},
					},
				},
			},
			list: func() ([]extensions.DaemonSet, error) {
				pod := &api.Pod{
					ObjectMeta: api.ObjectMeta{
						Name:      "pod1",
						Labels:    map[string]string{"foo": "bar"},
						Namespace: "ns",
					},
				}
				return lister.GetPodDaemonSets(pod)
			},
			outDaemonSetNames: sets.NewString("bar"),
		},
	}
	for _, c := range testCases {
		for _, r := range c.inDSs {
			store.Add(r)
		}

		daemonSets, err := c.list()
		if err != nil && c.expectErr {
			continue
		} else if c.expectErr {
			t.Error("Expected error, got none")
			continue
		} else if err != nil {
			t.Errorf("Unexpected error %#v", err)
			continue
		}
		daemonSetNames := make([]string, len(daemonSets))
		for ix := range daemonSets {
			daemonSetNames[ix] = daemonSets[ix].Name
		}
		if !c.outDaemonSetNames.HasAll(daemonSetNames...) || len(daemonSetNames) != len(c.outDaemonSetNames) {
			t.Errorf("Unexpected got controllers %+v expected %+v", daemonSetNames, c.outDaemonSetNames)
		}
	}
}

func TestStoreToPodLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	ids := []string{"foo", "bar", "baz"}
	for _, id := range ids {
		store.Add(&api.Pod{
			ObjectMeta: api.ObjectMeta{
				Name:   id,
				Labels: map[string]string{"name": id},
			},
		})
	}
	spl := StoreToPodLister{store}

	for _, id := range ids {
		got, err := spl.List(labels.Set{"name": id}.AsSelector())
		if err != nil {
			t.Errorf("Unexpected error: %v", err)
			continue
		}
		if e, a := 1, len(got); e != a {
			t.Errorf("Expected %v, got %v", e, a)
			continue
		}
		if e, a := id, got[0].Name; e != a {
			t.Errorf("Expected %v, got %v", e, a)
			continue
		}

		exists, err := spl.Exists(&api.Pod{ObjectMeta: api.ObjectMeta{Name: id}})
		if err != nil {
			t.Errorf("unexpected error: %v", err)
		}
		if !exists {
			t.Errorf("exists returned false for %v", id)
		}
	}

	exists, err := spl.Exists(&api.Pod{ObjectMeta: api.ObjectMeta{Name: "qux"}})
	if err != nil {
		t.Errorf("unexpected error: %v", err)
	}
	if exists {
		t.Error("Unexpected pod exists")
	}
}

func TestStoreToServiceLister(t *testing.T) {
	store := NewStore(MetaNamespaceKeyFunc)
	store.Add(&api.Service{
		ObjectMeta: api.ObjectMeta{Name: "foo"},
		Spec: api.ServiceSpec{
			Selector: map[string]string{},
		},
	})
	store.Add(&api.Service{ObjectMeta: api.ObjectMeta{Name: "bar"}})
	ssl := StoreToServiceLister{store}

	pod := &api.Pod{
		ObjectMeta: api.ObjectMeta{
			Name:   "foopod",
			Labels: map[string]string{"role": "foo"},
		},
	}

	services, err := ssl.GetPodServices(pod)
	if err != nil {
		t.Errorf("Unexpected error: %v", err)
	}

	if len(services) != 1 {
		t.Fatalf("Expected 1 service, got %v", len(services))
	}
	if e, a := "foo", services[0].Name; e != a {
		t.Errorf("Expected service %q, got %q", e, a)
	}
}
<|endoftext|>


<commit_before>package control

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"encoding/base64"

	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	"go.etcd.io/etcd/clientv3"
)

const (
	etcdDialTimeout    = 5 * time.Second
	k3sRuntimeEtcdPath = "/k3s/runtime"
	bootstrapTypeNone  = "none"
	bootstrapTypeRead  = "read"
	bootstrapTypeWrite = "write"
	bootstrapTypeFull  = "full"
)

type serverBootstrap struct {
	ServerCAData           string `json:"serverCAData,omitempty"`
	ServerCAKeyData        string `json:"serverCAKeyData,omitempty"`
	ClientCAData           string `json:"clientCAData,omitempty"`
	ClientCAKeyData        string `json:"clientCAKeyData,omitempty"`
	ServiceKeyData         string `json:"serviceKeyData,omitempty"`
	PasswdFileData         string `json:"passwdFileData,omitempty"`
	RequestHeaderCAData    string `json:"requestHeaderCAData,omitempty"`
	RequestHeaderCAKeyData string `json:"requestHeaderCAKeyData,omitempty"`
}

var validBootstrapTypes = map[string]bool{
	bootstrapTypeRead:  true,
	bootstrapTypeWrite: true,
	bootstrapTypeFull:  true,
}

func fetchBootstrapData(cfg *config.Control) error {
	if valid, err := checkBootstrapArgs(cfg, map[string]bool{
		bootstrapTypeFull: true,
		bootstrapTypeRead: true,
	}); !valid {
		if err != nil {
			logrus.Warnf("Not fetching bootstrap data: %v", err)
		}
		return nil
	}

	tlsConfig, err := genBootstrapTLSConfig(cfg)
	if err != nil {
		return err
	}

	endpoints := strings.Split(cfg.StorageEndpoint, ",")
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: etcdDialTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		return err
	}
	defer cli.Close()

	gr, err := cli.Get(context.TODO(), k3sRuntimeEtcdPath)
	if err != nil {
		return err
	}
	if len(gr.Kvs) == 0 {
		if cfg.BootstrapType != bootstrapTypeRead {
			return nil
		}
		return errors.New("Unable to read bootstrap data from server")
	}

	runtimeJSON, err := base64.URLEncoding.DecodeString(string(gr.Kvs[0].Value))
	if err != nil {
		return err
	}
	serverRuntime := &serverBootstrap{}
	if err := json.Unmarshal(runtimeJSON, serverRuntime); err != nil {
		return err
	}
	return writeRuntimeBootstrapData(cfg.Runtime, serverRuntime)
}

func storeBootstrapData(cfg *config.Control) error {
	if valid, err := checkBootstrapArgs(cfg, map[string]bool{
		bootstrapTypeFull:  true,
		bootstrapTypeWrite: true,
	}); !valid {
		if err != nil {
			logrus.Warnf("Not storing boostrap data: %v", err)
		}
		return nil
	}

	tlsConfig, err := genBootstrapTLSConfig(cfg)
	if err != nil {
		return err
	}

	endpoints := strings.Split(cfg.StorageEndpoint, ",")
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: etcdDialTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		return err
	}
	defer cli.Close()

	if cfg.BootstrapType != bootstrapTypeWrite {
		gr, err := cli.Get(context.TODO(), k3sRuntimeEtcdPath)
		if err != nil {
			return err
		}
		if len(gr.Kvs) > 0 && string(gr.Kvs[0].Value) != "" {
			return nil
		}
	}

	certData, err := readRuntimeBootstrapData(cfg.Runtime)
	if err != nil {
		return err
	}

	runtimeBase64 := base64.StdEncoding.EncodeToString(certData)
	_, err = cli.Put(context.TODO(), k3sRuntimeEtcdPath, runtimeBase64)
	if err != nil {
		return err
	}

	return nil
}

func checkBootstrapArgs(cfg *config.Control, accepted map[string]bool) (bool, error) {
	if cfg.BootstrapType == "" || cfg.BootstrapType == bootstrapTypeNone {
		return false, nil
	}
	if !validBootstrapTypes[cfg.BootstrapType] {
		return false, fmt.Errorf("unsupported bootstrap type [%s]", cfg.BootstrapType)
	}
	if cfg.StorageBackend != "etcd3" {
		return false, errors.New("bootstrap only supported with etcd3 as storage backend")
	}
	if !accepted[cfg.BootstrapType] {
		return false, nil
	}
	return true, nil
}

func genBootstrapTLSConfig(cfg *config.Control) (*tls.Config, error) {
	tlsConfig := &tls.Config{}
	if cfg.StorageCertFile != "" && cfg.StorageKeyFile != "" {
		certPem, err := ioutil.ReadFile(cfg.StorageCertFile)
		if err != nil {
			return nil, err
		}
		keyPem, err := ioutil.ReadFile(cfg.StorageKeyFile)
		if err != nil {
			return nil, err
		}
		tlsCert, err := tls.X509KeyPair(certPem, keyPem)
		if err != nil {
			return nil, err
		}
		tlsConfig.Certificates = []tls.Certificate{tlsCert}
	}
	if cfg.StorageCAFile != "" {
		caData, err := ioutil.ReadFile(cfg.StorageCAFile)
		if err != nil {
			return nil, err
		}
		certPool := x509.NewCertPool()
		certPool.AppendCertsFromPEM(caData)
		tlsConfig.RootCAs = certPool
	}
	return tlsConfig, nil
}

func readRuntimeBootstrapData(runtime *config.ControlRuntime) ([]byte, error) {
	serverBootstrapFiles := map[string]string{
		runtime.ServerCA:           "",
		runtime.ServerCAKey:        "",
		runtime.ClientCA:           "",
		runtime.ClientCAKey:        "",
		runtime.ServiceKey:         "",
		runtime.PasswdFile:         "",
		runtime.RequestHeaderCA:    "",
		runtime.RequestHeaderCAKey: "",
	}
	for k := range serverBootstrapFiles {
		data, err := ioutil.ReadFile(k)
		if err != nil {
			return nil, err
		}
		serverBootstrapFiles[k] = string(data)
	}
	serverBootstrapFileData := &serverBootstrap{
		ServerCAData:           serverBootstrapFiles[runtime.ServerCA],
		ServerCAKeyData:        serverBootstrapFiles[runtime.ServerCAKey],
		ClientCAData:           serverBootstrapFiles[runtime.ClientCA],
		ClientCAKeyData:        serverBootstrapFiles[runtime.ClientCAKey],
		ServiceKeyData:         serverBootstrapFiles[runtime.ServiceKey],
		PasswdFileData:         serverBootstrapFiles[runtime.PasswdFile],
		RequestHeaderCAData:    serverBootstrapFiles[runtime.RequestHeaderCA],
		RequestHeaderCAKeyData: serverBootstrapFiles[runtime.RequestHeaderCAKey],
	}
	return json.Marshal(serverBootstrapFileData)
}

func writeRuntimeBootstrapData(runtime *config.ControlRuntime, runtimeData *serverBootstrap) error {
	runtimePathValue := map[string]string{
		runtime.ServerCA:           runtimeData.ServerCAData,
		runtime.ServerCAKey:        runtimeData.ServerCAKeyData,
		runtime.ClientCA:           runtimeData.ClientCAData,
		runtime.ClientCAKey:        runtimeData.ClientCAKeyData,
		runtime.ServiceKey:         runtimeData.ServiceKeyData,
		runtime.PasswdFile:         runtimeData.PasswdFileData,
		runtime.RequestHeaderCA:    runtimeData.RequestHeaderCAData,
		runtime.RequestHeaderCAKey: runtimeData.RequestHeaderCAKeyData,
	}
	for k, v := range runtimePathValue {
		if _, err := os.Stat(k); os.IsNotExist(err) {
			if err := ioutil.WriteFile(k, []byte(v), 600); err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>Add more logs for bootstrap<commit_after>package control

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"time"

	"encoding/base64"

	"github.com/rancher/k3s/pkg/daemons/config"
	"github.com/sirupsen/logrus"
	"go.etcd.io/etcd/clientv3"
)

const (
	etcdDialTimeout    = 5 * time.Second
	k3sRuntimeEtcdPath = "/k3s/runtime"
	bootstrapTypeNone  = "none"
	bootstrapTypeRead  = "read"
	bootstrapTypeWrite = "write"
	bootstrapTypeFull  = "full"
)

type serverBootstrap struct {
	ServerCAData           string `json:"serverCAData,omitempty"`
	ServerCAKeyData        string `json:"serverCAKeyData,omitempty"`
	ClientCAData           string `json:"clientCAData,omitempty"`
	ClientCAKeyData        string `json:"clientCAKeyData,omitempty"`
	ServiceKeyData         string `json:"serviceKeyData,omitempty"`
	PasswdFileData         string `json:"passwdFileData,omitempty"`
	RequestHeaderCAData    string `json:"requestHeaderCAData,omitempty"`
	RequestHeaderCAKeyData string `json:"requestHeaderCAKeyData,omitempty"`
}

var validBootstrapTypes = map[string]bool{
	bootstrapTypeRead:  true,
	bootstrapTypeWrite: true,
	bootstrapTypeFull:  true,
}

func fetchBootstrapData(cfg *config.Control) error {
	if valid, err := checkBootstrapArgs(cfg, map[string]bool{
		bootstrapTypeFull: true,
		bootstrapTypeRead: true,
	}); !valid {
		if err != nil {
			logrus.Warnf("Not fetching bootstrap data: %v", err)
		}
		return nil
	}

	tlsConfig, err := genBootstrapTLSConfig(cfg)
	if err != nil {
		return err
	}

	endpoints := strings.Split(cfg.StorageEndpoint, ",")
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: etcdDialTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		return err
	}
	defer cli.Close()

	logrus.Info("Fetching bootstrap data from etcd")
	gr, err := cli.Get(context.TODO(), k3sRuntimeEtcdPath)
	if err != nil {
		return err
	}
	if len(gr.Kvs) == 0 {
		if cfg.BootstrapType != bootstrapTypeRead {
			return nil
		}
		return errors.New("Unable to read bootstrap data from server")
	}

	runtimeJSON, err := base64.URLEncoding.DecodeString(string(gr.Kvs[0].Value))
	if err != nil {
		return err
	}
	serverRuntime := &serverBootstrap{}
	if err := json.Unmarshal(runtimeJSON, serverRuntime); err != nil {
		return err
	}
	return writeRuntimeBootstrapData(cfg.Runtime, serverRuntime)
}

func storeBootstrapData(cfg *config.Control) error {
	if valid, err := checkBootstrapArgs(cfg, map[string]bool{
		bootstrapTypeFull:  true,
		bootstrapTypeWrite: true,
	}); !valid {
		if err != nil {
			logrus.Warnf("Not storing boostrap data: %v", err)
		}
		return nil
	}

	tlsConfig, err := genBootstrapTLSConfig(cfg)
	if err != nil {
		return err
	}

	endpoints := strings.Split(cfg.StorageEndpoint, ",")
	cli, err := clientv3.New(clientv3.Config{
		Endpoints:   endpoints,
		DialTimeout: etcdDialTimeout,
		TLS:         tlsConfig,
	})
	if err != nil {
		return err
	}
	defer cli.Close()

	if cfg.BootstrapType != bootstrapTypeWrite {
		gr, err := cli.Get(context.TODO(), k3sRuntimeEtcdPath)
		if err != nil {
			return err
		}
		if len(gr.Kvs) > 0 && string(gr.Kvs[0].Value) != "" {
			return nil
		}
	}

	certData, err := readRuntimeBootstrapData(cfg.Runtime)
	if err != nil {
		return err
	}

	logrus.Info("Storing bootstrap data to etcd")
	runtimeBase64 := base64.StdEncoding.EncodeToString(certData)
	_, err = cli.Put(context.TODO(), k3sRuntimeEtcdPath, runtimeBase64)
	if err != nil {
		return err
	}

	return nil
}

func checkBootstrapArgs(cfg *config.Control, accepted map[string]bool) (bool, error) {
	if cfg.BootstrapType == "" || cfg.BootstrapType == bootstrapTypeNone {
		return false, nil
	}
	if !validBootstrapTypes[cfg.BootstrapType] {
		return false, fmt.Errorf("unsupported bootstrap type [%s]", cfg.BootstrapType)
	}
	if cfg.StorageBackend != "etcd3" {
		return false, errors.New("bootstrap only supported with etcd3 as storage backend")
	}
	if !accepted[cfg.BootstrapType] {
		return false, nil
	}
	return true, nil
}

func genBootstrapTLSConfig(cfg *config.Control) (*tls.Config, error) {
	tlsConfig := &tls.Config{}
	if cfg.StorageCertFile != "" && cfg.StorageKeyFile != "" {
		certPem, err := ioutil.ReadFile(cfg.StorageCertFile)
		if err != nil {
			return nil, err
		}
		keyPem, err := ioutil.ReadFile(cfg.StorageKeyFile)
		if err != nil {
			return nil, err
		}
		tlsCert, err := tls.X509KeyPair(certPem, keyPem)
		if err != nil {
			return nil, err
		}
		tlsConfig.Certificates = []tls.Certificate{tlsCert}
	}
	if cfg.StorageCAFile != "" {
		caData, err := ioutil.ReadFile(cfg.StorageCAFile)
		if err != nil {
			return nil, err
		}
		certPool := x509.NewCertPool()
		certPool.AppendCertsFromPEM(caData)
		tlsConfig.RootCAs = certPool
	}
	return tlsConfig, nil
}

func readRuntimeBootstrapData(runtime *config.ControlRuntime) ([]byte, error) {
	serverBootstrapFiles := map[string]string{
		runtime.ServerCA:           "",
		runtime.ServerCAKey:        "",
		runtime.ClientCA:           "",
		runtime.ClientCAKey:        "",
		runtime.ServiceKey:         "",
		runtime.PasswdFile:         "",
		runtime.RequestHeaderCA:    "",
		runtime.RequestHeaderCAKey: "",
	}
	for k := range serverBootstrapFiles {
		data, err := ioutil.ReadFile(k)
		if err != nil {
			return nil, err
		}
		serverBootstrapFiles[k] = string(data)
	}
	serverBootstrapFileData := &serverBootstrap{
		ServerCAData:           serverBootstrapFiles[runtime.ServerCA],
		ServerCAKeyData:        serverBootstrapFiles[runtime.ServerCAKey],
		ClientCAData:           serverBootstrapFiles[runtime.ClientCA],
		ClientCAKeyData:        serverBootstrapFiles[runtime.ClientCAKey],
		ServiceKeyData:         serverBootstrapFiles[runtime.ServiceKey],
		PasswdFileData:         serverBootstrapFiles[runtime.PasswdFile],
		RequestHeaderCAData:    serverBootstrapFiles[runtime.RequestHeaderCA],
		RequestHeaderCAKeyData: serverBootstrapFiles[runtime.RequestHeaderCAKey],
	}
	return json.Marshal(serverBootstrapFileData)
}

func writeRuntimeBootstrapData(runtime *config.ControlRuntime, runtimeData *serverBootstrap) error {
	runtimePathValue := map[string]string{
		runtime.ServerCA:           runtimeData.ServerCAData,
		runtime.ServerCAKey:        runtimeData.ServerCAKeyData,
		runtime.ClientCA:           runtimeData.ClientCAData,
		runtime.ClientCAKey:        runtimeData.ClientCAKeyData,
		runtime.ServiceKey:         runtimeData.ServiceKeyData,
		runtime.PasswdFile:         runtimeData.PasswdFileData,
		runtime.RequestHeaderCA:    runtimeData.RequestHeaderCAData,
		runtime.RequestHeaderCAKey: runtimeData.RequestHeaderCAKeyData,
	}
	for k, v := range runtimePathValue {
		if _, err := os.Stat(k); os.IsNotExist(err) {
			if err := ioutil.WriteFile(k, []byte(v), 600); err != nil {
				return err
			}
		}
	}
	return nil
}
<|endoftext|>