<commit_before>package world

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fake_cc"
	"github.com/cloudfoundry-incubator/inigo/loggregator_runner"
	wardenrunner "github.com/cloudfoundry-incubator/warden-linux/integration/runner"
	gorouterconfig "github.com/cloudfoundry/gorouter/config"
	"github.com/fraenkel/candiedyaml"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type BuiltExecutables map[string]string
type BuiltCircuses map[string]string

const CircusZipFilename = "some-circus.tgz"

type BuiltArtifacts struct {
	Executables BuiltExecutables
	Circuses    BuiltCircuses
}

type ComponentAddresses struct {
	NATS           string
	Etcd           string
	EtcdPeer       string
	LoggregatorIn  string
	LoggregatorOut string
	Executor       string
	Rep            string
	FakeCC         string
	FileServer     string
	Router         string
	TPS            string
	WardenLinux    string
}

type ComponentMaker struct {
	Artifacts BuiltArtifacts
	Addresses ComponentAddresses

	ExternalAddress string

	Stack string

	WardenBinPath    string
	WardenRootFSPath string
	WardenGraphPath  string
}

func (maker ComponentMaker) NATS(argv ...string) ifrit.Runner {
	host, port, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "gnatsd",
		AnsiColorCode:     "30m",
		StartCheck:        "gnatsd is ready",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"gnatsd",
			append([]string{
				"--addr", host,
				"--port", port,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Etcd(argv ...string) ifrit.Runner {
	nodeName := fmt.Sprintf("etcd_%d", ginkgo.GinkgoParallelNode())
	dataDir := path.Join(os.TempDir(), nodeName)

	return &ginkgomon.Runner{
		Name:              "etcd",
		AnsiColorCode:     "31m",
		StartCheck:        "leader changed",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"etcd",
			append([]string{
				"-data-dir", dataDir,
				"-addr", maker.Addresses.Etcd,
				"-peer-addr", maker.Addresses.EtcdPeer,
				"-name", nodeName,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(dataDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) WardenLinux(argv ...string) *wardenrunner.Runner {
	return wardenrunner.New(
		"tcp",
		maker.Addresses.WardenLinux,
		maker.Artifacts.Executables["warden-linux"],
		maker.WardenBinPath,
		maker.WardenRootFSPath,
		argv...,
	)
}

func (maker ComponentMaker) Executor(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "executor",
		AnsiColorCode:     "91m",
		StartCheck:        "executor.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["exec"],
			append([]string{
				"-listenAddr", maker.Addresses.Executor,
				"-wardenNetwork", "tcp",
				"-wardenAddr", maker.Addresses.WardenLinux,
				"-loggregatorServer", maker.Addresses.LoggregatorIn,
				"-loggregatorSecret", "loggregator-secret",
				"-containerMaxCpuShares", "1024",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Rep(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "rep",
		AnsiColorCode:     "92m",
		StartCheck:        "rep.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["rep"],
			append(
				[]string{
					"-stack", maker.Stack,
					"-lrpHost", maker.ExternalAddress,
					"-listenAddr", maker.Addresses.Rep,
					"-etcdCluster", "http://" + maker.Addresses.Etcd,
					"-natsAddresses", maker.Addresses.NATS,
					"-executorID", "the-executor-id-" + strconv.Itoa(ginkgo.GinkgoParallelNode()),
					"-executorURL", "http://" + maker.Addresses.Executor,
					"-heartbeatInterval", "1s",
				},
				argv...,
			)...,
		),
	}
}

func (maker ComponentMaker) Converger(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "converger",
		AnsiColorCode:     "93m",
		StartCheck:        "converger.started",
		StartCheckTimeout: 5 * time.Second,

		Command: exec.Command(
			maker.Artifacts.Executables["converger"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Auctioneer(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "auctioneer",
		AnsiColorCode: "94m",
		StartCheck:    "auctioneer.started",
		Command: exec.Command(
			maker.Artifacts.Executables["auctioneer"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,

				// we limit this to prevent overwhelming numbers of auctioneer logs.  it
				// should not impact the behavior of the tests.
				"-maxRounds", "3",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) RouteEmitter(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "route-emitter",
		AnsiColorCode: "95m",
		StartCheck:    "route-emitter.started",
		Command: exec.Command(
			maker.Artifacts.Executables["route-emitter"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) TPS(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "tps",
		AnsiColorCode: "96m",
		StartCheck:    "tps.started",
		Command: exec.Command(
			maker.Artifacts.Executables["tps"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-listenAddr", maker.Addresses.TPS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) NsyncListener(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "nsync-listener",
		AnsiColorCode: "97m",
		StartCheck:    "nsync.listener.started",
		Command: exec.Command(
			maker.Artifacts.Executables["nsync-listener"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-repAddrRelativeToExecutor", maker.Addresses.Rep,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) FileServer(argv ...string) (ifrit.Runner, string) {
	servedFilesDir, err := ioutil.TempDir("", "file-server-files")
	Ω(err).ShouldNot(HaveOccurred())

	host, port, err := net.SplitHostPort(maker.Addresses.FileServer)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "file-server",
		AnsiColorCode: "90m",
		StartCheck:    "file-server.ready",
		Command: exec.Command(
			maker.Artifacts.Executables["file-server"],
			append([]string{
				"-address", host,
				"-port", port,
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-ccAddress", "http://" + maker.Addresses.FakeCC,
				"-ccJobPollingInterval", "100ms",
				"-ccUsername", fake_cc.CC_USERNAME,
				"-ccPassword", fake_cc.CC_PASSWORD,
				"-staticDirectory", servedFilesDir,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(servedFilesDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}, servedFilesDir
}

func (maker ComponentMaker) Router() ifrit.Runner {
	_, routerPort, err := net.SplitHostPort(maker.Addresses.Router)
	Ω(err).ShouldNot(HaveOccurred())

	routerPortInt, err := strconv.Atoi(routerPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	routerConfig := &gorouterconfig.Config{
		Port: uint16(routerPortInt),

		PruneStaleDropletsIntervalInSeconds: 5,
		DropletStaleThresholdInSeconds:      10,
		PublishActiveAppsIntervalInSeconds:  0,
		StartResponseDelayIntervalInSeconds: 1,

		Nats: []gorouterconfig.NatsConfig{
			{
				Host: natsHost,
				Port: uint16(natsPortInt),
			},
		},
		Logging: gorouterconfig.LoggingConfig{
			File:  "/dev/stdout",
			Level: "info",
		},
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "router-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = candiedyaml.NewEncoder(configFile).Encode(routerConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "router",
		AnsiColorCode:     "32m",
		StartCheck:        "router.started",
		StartCheckTimeout: 5 * time.Second, // it waits 1 second before listening. yep.
		Command: exec.Command(
			maker.Artifacts.Executables["router"],
			"-c", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) Loggregator() ifrit.Runner {
	_, inPort, err := net.SplitHostPort(maker.Addresses.LoggregatorIn)
	Ω(err).ShouldNot(HaveOccurred())

	_, outPort, err := net.SplitHostPort(maker.Addresses.LoggregatorOut)
	Ω(err).ShouldNot(HaveOccurred())

	inPortInt, err := strconv.Atoi(inPort)
	Ω(err).ShouldNot(HaveOccurred())

	outPortInt, err := strconv.Atoi(outPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	loggregatorConfig := loggregator_runner.Config{
		IncomingPort:           inPortInt,
		OutgoingPort:           outPortInt,
		MaxRetainedLogMessages: 1000,
		SharedSecret:           "loggregator-secret",
		NatsHost:               natsHost,
		NatsPort:               natsPortInt,
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "loggregator-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = json.NewEncoder(configFile).Encode(loggregatorConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "loggregator",
		AnsiColorCode: "33m",
		StartCheck:    "Listening on port",
		Command: exec.Command(
			maker.Artifacts.Executables["loggregator"],
			"-config", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) FakeCC() *fake_cc.FakeCC {
	return fake_cc.New(maker.Addresses.FakeCC)
}

func (maker ComponentMaker) Stager(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "stager",
		AnsiColorCode: "94m",
		StartCheck:    "Listening for staging requests!",
		Command: exec.Command(
			maker.Artifacts.Executables["stager"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}
<commit_msg>bump auction nats timeout to 5 seconds<commit_after>package world

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"time"

	"github.com/cloudfoundry-incubator/inigo/fake_cc"
	"github.com/cloudfoundry-incubator/inigo/loggregator_runner"
	wardenrunner "github.com/cloudfoundry-incubator/warden-linux/integration/runner"
	gorouterconfig "github.com/cloudfoundry/gorouter/config"
	"github.com/fraenkel/candiedyaml"
	"github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
)

type BuiltExecutables map[string]string
type BuiltCircuses map[string]string

const CircusZipFilename = "some-circus.tgz"

type BuiltArtifacts struct {
	Executables BuiltExecutables
	Circuses    BuiltCircuses
}

type ComponentAddresses struct {
	NATS           string
	Etcd           string
	EtcdPeer       string
	LoggregatorIn  string
	LoggregatorOut string
	Executor       string
	Rep            string
	FakeCC         string
	FileServer     string
	Router         string
	TPS            string
	WardenLinux    string
}

type ComponentMaker struct {
	Artifacts BuiltArtifacts
	Addresses ComponentAddresses

	ExternalAddress string

	Stack string

	WardenBinPath    string
	WardenRootFSPath string
	WardenGraphPath  string
}

func (maker ComponentMaker) NATS(argv ...string) ifrit.Runner {
	host, port, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "gnatsd",
		AnsiColorCode:     "30m",
		StartCheck:        "gnatsd is ready",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"gnatsd",
			append([]string{
				"--addr", host,
				"--port", port,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Etcd(argv ...string) ifrit.Runner {
	nodeName := fmt.Sprintf("etcd_%d", ginkgo.GinkgoParallelNode())
	dataDir := path.Join(os.TempDir(), nodeName)

	return &ginkgomon.Runner{
		Name:              "etcd",
		AnsiColorCode:     "31m",
		StartCheck:        "leader changed",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			"etcd",
			append([]string{
				"-data-dir", dataDir,
				"-addr", maker.Addresses.Etcd,
				"-peer-addr", maker.Addresses.EtcdPeer,
				"-name", nodeName,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(dataDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) WardenLinux(argv ...string) *wardenrunner.Runner {
	return wardenrunner.New(
		"tcp",
		maker.Addresses.WardenLinux,
		maker.Artifacts.Executables["warden-linux"],
		maker.WardenBinPath,
		maker.WardenRootFSPath,
		argv...,
	)
}

func (maker ComponentMaker) Executor(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "executor",
		AnsiColorCode:     "91m",
		StartCheck:        "executor.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["exec"],
			append([]string{
				"-listenAddr", maker.Addresses.Executor,
				"-wardenNetwork", "tcp",
				"-wardenAddr", maker.Addresses.WardenLinux,
				"-loggregatorServer", maker.Addresses.LoggregatorIn,
				"-loggregatorSecret", "loggregator-secret",
				"-containerMaxCpuShares", "1024",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Rep(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "rep",
		AnsiColorCode:     "92m",
		StartCheck:        "rep.started",
		StartCheckTimeout: 5 * time.Second,
		Command: exec.Command(
			maker.Artifacts.Executables["rep"],
			append(
				[]string{
					"-stack", maker.Stack,
					"-lrpHost", maker.ExternalAddress,
					"-listenAddr", maker.Addresses.Rep,
					"-etcdCluster", "http://" + maker.Addresses.Etcd,
					"-natsAddresses", maker.Addresses.NATS,
					"-executorID", "the-executor-id-" + strconv.Itoa(ginkgo.GinkgoParallelNode()),
					"-executorURL", "http://" + maker.Addresses.Executor,
					"-heartbeatInterval", "1s",
				},
				argv...,
			)...,
		),
	}
}

func (maker ComponentMaker) Converger(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:              "converger",
		AnsiColorCode:     "93m",
		StartCheck:        "converger.started",
		StartCheckTimeout: 5 * time.Second,

		Command: exec.Command(
			maker.Artifacts.Executables["converger"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) Auctioneer(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "auctioneer",
		AnsiColorCode: "94m",
		StartCheck:    "auctioneer.started",
		Command: exec.Command(
			maker.Artifacts.Executables["auctioneer"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,

				// inigo runs everything on the same machine, so there will be more
				// load; this timeout is a bit sensitive to that.
				"-natsAuctionTimeout", "5s",

				// we limit this to prevent overwhelming numbers of auctioneer logs.  it
				// should not impact the behavior of the tests.
				"-maxRounds", "3",
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) RouteEmitter(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "route-emitter",
		AnsiColorCode: "95m",
		StartCheck:    "route-emitter.started",
		Command: exec.Command(
			maker.Artifacts.Executables["route-emitter"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) TPS(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "tps",
		AnsiColorCode: "96m",
		StartCheck:    "tps.started",
		Command: exec.Command(
			maker.Artifacts.Executables["tps"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-listenAddr", maker.Addresses.TPS,
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) NsyncListener(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "nsync-listener",
		AnsiColorCode: "97m",
		StartCheck:    "nsync.listener.started",
		Command: exec.Command(
			maker.Artifacts.Executables["nsync-listener"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-repAddrRelativeToExecutor", maker.Addresses.Rep,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}

func (maker ComponentMaker) FileServer(argv ...string) (ifrit.Runner, string) {
	servedFilesDir, err := ioutil.TempDir("", "file-server-files")
	Ω(err).ShouldNot(HaveOccurred())

	host, port, err := net.SplitHostPort(maker.Addresses.FileServer)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "file-server",
		AnsiColorCode: "90m",
		StartCheck:    "file-server.ready",
		Command: exec.Command(
			maker.Artifacts.Executables["file-server"],
			append([]string{
				"-address", host,
				"-port", port,
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-ccAddress", "http://" + maker.Addresses.FakeCC,
				"-ccJobPollingInterval", "100ms",
				"-ccUsername", fake_cc.CC_USERNAME,
				"-ccPassword", fake_cc.CC_PASSWORD,
				"-staticDirectory", servedFilesDir,
			}, argv...)...,
		),
		Cleanup: func() {
			err := os.RemoveAll(servedFilesDir)
			Ω(err).ShouldNot(HaveOccurred())
		},
	}, servedFilesDir
}

func (maker ComponentMaker) Router() ifrit.Runner {
	_, routerPort, err := net.SplitHostPort(maker.Addresses.Router)
	Ω(err).ShouldNot(HaveOccurred())

	routerPortInt, err := strconv.Atoi(routerPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	routerConfig := &gorouterconfig.Config{
		Port: uint16(routerPortInt),

		PruneStaleDropletsIntervalInSeconds: 5,
		DropletStaleThresholdInSeconds:      10,
		PublishActiveAppsIntervalInSeconds:  0,
		StartResponseDelayIntervalInSeconds: 1,

		Nats: []gorouterconfig.NatsConfig{
			{
				Host: natsHost,
				Port: uint16(natsPortInt),
			},
		},
		Logging: gorouterconfig.LoggingConfig{
			File:  "/dev/stdout",
			Level: "info",
		},
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "router-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = candiedyaml.NewEncoder(configFile).Encode(routerConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:              "router",
		AnsiColorCode:     "32m",
		StartCheck:        "router.started",
		StartCheckTimeout: 5 * time.Second, // it waits 1 second before listening. yep.
		Command: exec.Command(
			maker.Artifacts.Executables["router"],
			"-c", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) Loggregator() ifrit.Runner {
	_, inPort, err := net.SplitHostPort(maker.Addresses.LoggregatorIn)
	Ω(err).ShouldNot(HaveOccurred())

	_, outPort, err := net.SplitHostPort(maker.Addresses.LoggregatorOut)
	Ω(err).ShouldNot(HaveOccurred())

	inPortInt, err := strconv.Atoi(inPort)
	Ω(err).ShouldNot(HaveOccurred())

	outPortInt, err := strconv.Atoi(outPort)
	Ω(err).ShouldNot(HaveOccurred())

	natsHost, natsPort, err := net.SplitHostPort(maker.Addresses.NATS)
	Ω(err).ShouldNot(HaveOccurred())

	natsPortInt, err := strconv.Atoi(natsPort)
	Ω(err).ShouldNot(HaveOccurred())

	loggregatorConfig := loggregator_runner.Config{
		IncomingPort:           inPortInt,
		OutgoingPort:           outPortInt,
		MaxRetainedLogMessages: 1000,
		SharedSecret:           "loggregator-secret",
		NatsHost:               natsHost,
		NatsPort:               natsPortInt,
	}

	configFile, err := ioutil.TempFile(os.TempDir(), "loggregator-config")
	Ω(err).ShouldNot(HaveOccurred())

	defer configFile.Close()

	err = json.NewEncoder(configFile).Encode(loggregatorConfig)
	Ω(err).ShouldNot(HaveOccurred())

	return &ginkgomon.Runner{
		Name:          "loggregator",
		AnsiColorCode: "33m",
		StartCheck:    "Listening on port",
		Command: exec.Command(
			maker.Artifacts.Executables["loggregator"],
			"-config", configFile.Name(),
		),
		Cleanup: func() {
			err := os.Remove(configFile.Name())
			Ω(err).ShouldNot(HaveOccurred())
		},
	}
}

func (maker ComponentMaker) FakeCC() *fake_cc.FakeCC {
	return fake_cc.New(maker.Addresses.FakeCC)
}

func (maker ComponentMaker) Stager(argv ...string) ifrit.Runner {
	return &ginkgomon.Runner{
		Name:          "stager",
		AnsiColorCode: "94m",
		StartCheck:    "Listening for staging requests!",
		Command: exec.Command(
			maker.Artifacts.Executables["stager"],
			append([]string{
				"-etcdCluster", "http://" + maker.Addresses.Etcd,
				"-natsAddresses", maker.Addresses.NATS,
				"-circuses", fmt.Sprintf(`{"%s": "%s"}`, maker.Stack, CircusZipFilename),
			}, argv...)...,
		),
	}
}
<|endoftext|>


<commit_before>package tail

import (
	"os"

	"github.com/hpcloud/tail"
)

// Follower describes an object that continuously emits a stream of lines
type Follower interface {
	Lines() chan *tail.Line
	OnError(func(error))
}

type followerImpl struct {
	filename string
	t        *tail.Tail
}

// NewFollower creates a new Follower instance for a given file (given by name)
func NewFollower(filename string) (Follower, error) {
	f := &followerImpl{
		filename: filename,
	}

	if err := f.start(); err != nil {
		return nil, err
	}

	return f, nil
}

func (f *followerImpl) start() error {
	seekInfo := tail.SeekInfo{Offset: 0, Whence: os.SEEK_END}
	t, err := tail.TailFile(f.filename, tail.Config{
		Follow:   true,
		ReOpen:   true,
		Poll:     true,
		Location: &seekInfo,
	})

	if err != nil {
		return err
	}

	f.t = t
	return nil
}

func (f *followerImpl) OnError(cb func(error)) {
	go func() {
		err := f.t.Wait()
		if err != nil {
			cb(err)
		}
	}()
}

func (f *followerImpl) Lines() chan *tail.Line {
	return f.t.Lines
}
<commit_msg>Only seek from end if file existed before<commit_after>package tail

import (
	"os"

	"github.com/hpcloud/tail"
)

// Follower describes an object that continuously emits a stream of lines
type Follower interface {
	Lines() chan *tail.Line
	OnError(func(error))
}

type followerImpl struct {
	filename string
	t        *tail.Tail
}

// NewFollower creates a new Follower instance for a given file (given by name)
func NewFollower(filename string) (Follower, error) {
	f := &followerImpl{
		filename: filename,
	}

	if err := f.start(); err != nil {
		return nil, err
	}

	return f, nil
}

func (f *followerImpl) start() error {
	var seekInfo *tail.SeekInfo

	_, err := os.Stat(f.filename)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	} else {
		seekInfo = &tail.SeekInfo{Offset: 0, Whence: os.SEEK_END}
	}

	t, err := tail.TailFile(f.filename, tail.Config{
		Follow:   true,
		ReOpen:   true,
		Poll:     true,
		Location: seekInfo,
	})

	if err != nil {
		return err
	}

	f.t = t
	return nil
}

func (f *followerImpl) OnError(cb func(error)) {
	go func() {
		err := f.t.Wait()
		if err != nil {
			cb(err)
		}
	}()
}

func (f *followerImpl) Lines() chan *tail.Line {
	return f.t.Lines
}
<|endoftext|>


<commit_before>package identify

import (
	"testing"
	"time"

	ma "github.com/multiformats/go-multiaddr"
)

// TestObsAddrSet
func TestObsAddrSet(t *testing.T) {
	m := func(s string) ma.Multiaddr {
		m, err := ma.NewMultiaddr(s)
		if err != nil {
			t.Error(err)
		}
		return m
	}

	addrsMarch := func(a, b []ma.Multiaddr) bool {
		if len(a) != len(b) {
			return false
		}

		for _, aa := range a {
			found := false
			for _, bb := range b {
				if aa.Equal(bb) {
					found = true
					break
				}
			}
			if !found {
				return false
			}
		}
		return true
	}

	a1 := m("/ip4/1.2.3.4/tcp/1231")
	a2 := m("/ip4/1.2.3.4/tcp/1232")
	a3 := m("/ip4/1.2.3.4/tcp/1233")
	a4 := m("/ip4/1.2.3.4/tcp/1234")
	a5 := m("/ip4/1.2.3.4/tcp/1235")

	b1 := m("/ip4/1.2.3.6/tcp/1236")
	b2 := m("/ip4/1.2.3.7/tcp/1237")
	b3 := m("/ip4/1.2.3.8/tcp/1237")
	b4 := m("/ip4/1.2.3.9/tcp/1237")
	b5 := m("/ip4/1.2.3.10/tcp/1237")

	oas := ObservedAddrSet{}

	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should be empty")
	}

	oas.Add(a1, a4)
	oas.Add(a2, a4)
	oas.Add(a3, a4)

	// these are all different so we should not yet get them.
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (once)")
	}

	// same observer, so should not yet get them.
	oas.Add(a1, a4)
	oas.Add(a2, a4)
	oas.Add(a3, a4)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (same obs)")
	}

	// different observer, but same observer group.
	oas.Add(a1, a5)
	oas.Add(a2, a5)
	oas.Add(a3, a5)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (same obs group)")
	}

	oas.Add(a1, b1)
	oas.Add(a1, b2)
	oas.Add(a1, b3)
	if !addrsMarch(oas.Addrs(), []ma.Multiaddr{a1}) {
		t.Error("addrs should only have a1")
	}

	oas.Add(a2, a5)
	oas.Add(a1, a5)
	oas.Add(a1, a5)
	oas.Add(a2, b1)
	oas.Add(a1, b1)
	oas.Add(a1, b1)
	oas.Add(a2, b2)
	oas.Add(a1, b2)
	oas.Add(a1, b2)
	oas.Add(a2, b4)
	oas.Add(a2, b5)
	if !addrsMarch(oas.Addrs(), []ma.Multiaddr{a1, a2}) {
		t.Error("addrs should only have a1, a2")
	}

	// change the timeout constant so we can time it out.
	oas.SetTTL(time.Millisecond * 200)
	<-time.After(time.Millisecond * 210)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should have timed out")
	}
}
<commit_msg>Fix test<commit_after>package identify

import (
	"testing"
	"time"

	net "github.com/libp2p/go-libp2p-net"
	ma "github.com/multiformats/go-multiaddr"
)

// TestObsAddrSet
func TestObsAddrSet(t *testing.T) {
	m := func(s string) ma.Multiaddr {
		m, err := ma.NewMultiaddr(s)
		if err != nil {
			t.Error(err)
		}
		return m
	}

	addrsMarch := func(a, b []ma.Multiaddr) bool {
		if len(a) != len(b) {
			return false
		}

		for _, aa := range a {
			found := false
			for _, bb := range b {
				if aa.Equal(bb) {
					found = true
					break
				}
			}
			if !found {
				return false
			}
		}
		return true
	}

	a1 := m("/ip4/1.2.3.4/tcp/1231")
	a2 := m("/ip4/1.2.3.4/tcp/1232")
	a3 := m("/ip4/1.2.3.4/tcp/1233")
	a4 := m("/ip4/1.2.3.4/tcp/1234")
	a5 := m("/ip4/1.2.3.4/tcp/1235")

	b1 := m("/ip4/1.2.3.6/tcp/1236")
	b2 := m("/ip4/1.2.3.7/tcp/1237")
	b3 := m("/ip4/1.2.3.8/tcp/1237")
	b4 := m("/ip4/1.2.3.9/tcp/1237")
	b5 := m("/ip4/1.2.3.10/tcp/1237")

	oas := &ObservedAddrSet{}

	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should be empty")
	}

	add := func(oas *ObservedAddrSet, observed, observer ma.Multiaddr) {
		dummyLocal := m("/ip4/127.0.0.1/tcp/10086")
		dummyDirection := net.DirOutbound

		oas.Add(observed, dummyLocal, observer, dummyDirection)
	}

	add(oas, a1, a4)
	add(oas, a2, a4)
	add(oas, a3, a4)

	// these are all different so we should not yet get them.
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (once)")
	}

	// same observer, so should not yet get them.
	add(oas, a1, a4)
	add(oas, a2, a4)
	add(oas, a3, a4)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (same obs)")
	}

	// different observer, but same observer group.
	add(oas, a1, a5)
	add(oas, a2, a5)
	add(oas, a3, a5)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should _still_ be empty (same obs group)")
	}

	add(oas, a1, b1)
	add(oas, a1, b2)
	add(oas, a1, b3)
	if !addrsMarch(oas.Addrs(), []ma.Multiaddr{a1}) {
		t.Error("addrs should only have a1")
	}

	add(oas, a2, a5)
	add(oas, a1, a5)
	add(oas, a1, a5)
	add(oas, a2, b1)
	add(oas, a1, b1)
	add(oas, a1, b1)
	add(oas, a2, b2)
	add(oas, a1, b2)
	add(oas, a1, b2)
	add(oas, a2, b4)
	add(oas, a2, b5)
	if !addrsMarch(oas.Addrs(), []ma.Multiaddr{a1, a2}) {
		t.Error("addrs should only have a1, a2")
	}

	// change the timeout constant so we can time it out.
	oas.SetTTL(time.Millisecond * 200)
	<-time.After(time.Millisecond * 210)
	if !addrsMarch(oas.Addrs(), nil) {
		t.Error("addrs should have timed out")
	}
}
<|endoftext|>


<commit_before>package workque

import (
	"log"
	"math/rand"
	"sync"
	"time"

	"github.com/Shopify/sarama"
)

var quit = make(chan *sync.WaitGroup, 1)

// Payload a Kafka consumer message
type Payload struct {
	Message sarama.ConsumerMessage
}

// Worker represents the worker that executes the job
type Worker struct {
	WorkerPool chan chan Job
	JobChannel chan Job
}

// NewWorker create new channel worker
func NewWorker(workerPool chan chan Job) Worker {
	return Worker{
		WorkerPool: workerPool,
		JobChannel: make(chan Job)}
}

// Start method starts the run loop for the worker,
// listening for a quit channel in case we need to stop it
func (w Worker) Start() {
	go func() {
		// register the current worker into the worker queue.
		w.WorkerPool <- w.JobChannel
		log.Println("Worker started...")

		for {
			select {
			case job := <-w.JobChannel:
				// we have received a work request.
				if err := job.Payload.InsertIntoDB(); err != nil {
					log.Printf("Error inserting into database: %s", err.Error())
				}

			case wg := <-quit:
				// we have received a signal to stop
				log.Println("Worker stopped")
				wg.Done()
				return
			}
		}
	}()
}

// Stop signals the worker to stop listening for work requests.
func (w Worker) Stop(wg *sync.WaitGroup) {
	go func() {
		log.Println("Stopping worker...")
		quit <- wg
	}()
}

// InsertIntoDB insert receive data into database
func (p *Payload) InsertIntoDB() error {
	log.Printf("worker process message - Topic: %s, Partition: %d, Offset: %d,\tMessageKey: %s,\tMessageValue: %s", p.Message.Topic, p.Message.Partition, p.Message.Offset, p.Message.Key, p.Message.Value)
	time.Sleep(time.Duration(rand.Intn(200)) * time.Millisecond)
	log.Println("work done")
	return nil
}
<commit_msg> insert into real database, and fix  issue worker stop after process<commit_after>package workque

import (
	"encoding/json"
	"log"
	"sync"
	"time"

	"github.com/gocql/gocql"
)

var quit = make(chan *sync.WaitGroup, 1)

// Data a MQTT payload model
type Data struct {
	Temperature string `json:"temperature"`
	Humidity    string `json:"humidity"`
}

// Worker represents the worker that executes the job
type Worker struct {
	WorkerPool chan chan Job
	JobChannel chan Job
}

// NewWorker create new channel worker
func NewWorker(workerPool chan chan Job) Worker {
	return Worker{
		WorkerPool: workerPool,
		JobChannel: make(chan Job)}
}

// Start method starts the run loop for the worker,
// listening for a quit channel in case we need to stop it
func (w Worker) Start() {
	go func() {
		log.Println("Worker started")

		for {
			// register the current worker into the worker queue.
			w.WorkerPool <- w.JobChannel

			select {
			case job := <-w.JobChannel:
				// we have received a work request.
				if err := job.InsertIntoDB(); err != nil {
					log.Printf("Error inserting into database: %s", err.Error())
				}

			case wg := <-quit:
				// we have received a signal to stop
				log.Println("Worker stopped")
				wg.Done()
				return
			}
		}
	}()
}

// Stop signals the worker to stop listening for work requests.
func (w Worker) Stop(wg *sync.WaitGroup) {
	go func() {
		log.Println("Stopping worker...")
		quit <- wg
	}()
}

// InsertIntoDB insert received data into database
func (j *Job) InsertIntoDB() error {
	log.Printf("worker process message - Topic: %s, Partition: %d, Offset: %d,\tMessageKey: %s,\tMessageValue: %s", j.Message.Topic, j.Message.Partition, j.Message.Offset, j.Message.Key, j.Message.Value)
	var data Data
	if err := json.Unmarshal(j.Message.Value, &data); err != nil {
		return err
	}

	if err := j.DB.Query(`INSERT INTO data (id, username, temperature, humidity, timestamp) VALUES (?, ?, ?, ?, ?)`,
		gocql.TimeUUID(), j.Message.Key, data.Temperature, data.Humidity, time.Now()).Exec(); err != nil {
		return err
	}
	log.Println("work done")
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2019 The LUCI 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 exe

import (
	"bytes"
	"compress/zlib"
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/golang/protobuf/proto"
	structpb "github.com/golang/protobuf/ptypes/struct"

	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/logdog/client/butlerlib/bootstrap"
	"go.chromium.org/luci/logdog/client/butlerlib/streamclient"
	"go.chromium.org/luci/luciexe"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func TestExe(t *testing.T) {
	t.Parallel()

	Convey(`test exe`, t, func() {
		client := streamclient.NewFake("test_namespace")
		ldClient := &bootstrap.Bootstrap{
			CoordinatorHost: "test.example.com",
			Project:         "test_project",
			Prefix:          "test_prefix",
			Namespace:       "test_namespace",
			Client:          client.Client,
		}
		var ldErr error
		bootstrapGet := func() (*bootstrap.Bootstrap, error) { return ldClient, ldErr }

		getBuilds := func(decompress bool) []*bbpb.Build {
			fakeData := client.GetFakeData()["test_namespace/build.proto"]
			if decompress {
				So(fakeData.GetFlags().ContentType, ShouldEqual, luciexe.BuildProtoZlibContentType)
			} else {
				So(fakeData.GetFlags().ContentType, ShouldEqual, luciexe.BuildProtoContentType)
			}

			dgs := fakeData.GetDatagrams()
			So(len(dgs), ShouldBeGreaterThanOrEqualTo, 1)

			ret := make([]*bbpb.Build, len(dgs))
			for i, dg := range dgs {
				ret[i] = &bbpb.Build{}

				var data []byte

				if decompress {
					r, err := zlib.NewReader(bytes.NewBufferString(dg))
					So(err, ShouldBeNil)
					data, err = ioutil.ReadAll(r)
					So(err, ShouldBeNil)
				} else {
					data = []byte(dg)
				}

				So(proto.Unmarshal(data, ret[i]), ShouldBeNil)
			}
			return ret
		}
		lastBuild := func() *bbpb.Build {
			builds := getBuilds(false)
			return builds[len(builds)-1]
		}

		args := []string{"fake_test_executable"}

		ctx := context.Background()

		Convey(`basic`, func() {
			Convey(`success`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.Status = bbpb.Status_SCHEDULED
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status: bbpb.Status_SUCCESS,
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`failure`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					return errors.New("bad stuff")
				})
				So(exitCode, ShouldEqual, 1)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_FAILURE,
					SummaryMarkdown: "Final error: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`infra failure`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					return errors.New("bad stuff", InfraErrorTag)
				})
				So(exitCode, ShouldEqual, 1)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "Final infra error: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`panic`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					panic(errors.New("bad stuff"))
				})
				So(exitCode, ShouldEqual, 2)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "Final panic: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`respect user program status`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.Status = bbpb.Status_INFRA_FAILURE
					build.SummaryMarkdown = "status set inside"
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "status set inside",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})
		})

		Convey(`send`, func() {
			exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
				build.SummaryMarkdown = "Hi. I did stuff."
				bs()
				return errors.New("oh no i failed")
			})
			So(exitCode, ShouldEqual, 1)
			builds := getBuilds(false)
			So(len(builds), ShouldEqual, 2)
			So(builds[0], ShouldResembleProto, &bbpb.Build{
				SummaryMarkdown: "Hi. I did stuff.",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
			So(builds[len(builds)-1], ShouldResembleProto, &bbpb.Build{
				Status:          bbpb.Status_FAILURE,
				SummaryMarkdown: "Hi. I did stuff.\n\nFinal error: oh no i failed",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
		})

		Convey(`send (zlib)`, func() {
			exitCode := runCtx(ctx, args, bootstrapGet, []Option{WithZlibCompression(5)}, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
				build.SummaryMarkdown = "Hi. I did stuff."
				bs()
				return errors.New("oh no i failed")
			})
			So(exitCode, ShouldEqual, 1)
			builds := getBuilds(true)
			So(len(builds), ShouldEqual, 2)
			So(builds[0], ShouldResembleProto, &bbpb.Build{
				SummaryMarkdown: "Hi. I did stuff.",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
			So(builds[len(builds)-1], ShouldResembleProto, &bbpb.Build{
				Status:          bbpb.Status_FAILURE,
				SummaryMarkdown: "Hi. I did stuff.\n\nFinal error: oh no i failed",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
		})

		Convey(`output`, func() {
			tdir, err := ioutil.TempDir("", "luciexe-exe-test")
			So(err, ShouldBeNil)
			defer os.RemoveAll(tdir)

			Convey(`binary`, func() {
				outFile := filepath.Join(tdir, "out.pb")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					err := WriteProperties(build.Output.Properties, map[string]interface{}{
						"some": "thing",
					})
					if err != nil {
						panic(err)
					}

					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{
							Fields: map[string]*structpb.Value{
								"some": {Kind: &structpb.Value_StringValue{
									StringValue: "thing",
								}},
							},
						},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"`\f\x82\x01\x13\n\x11\n\x0f\n\x04some\x12\a\x1a\x05thing\xa2\x01\x03Hi.")
			})

			Convey(`textpb`, func() {
				outFile := filepath.Join(tdir, "out.textpb")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"status: SUCCESS\nsummary_markdown: \"Hi.\"\noutput: <\n  properties: <\n  >\n>\n")
			})

			Convey(`jsonpb`, func() {
				outFile := filepath.Join(tdir, "out.json")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"{\n  \"status\": \"SUCCESS\",\n  \"summary_markdown\": \"Hi.\",\n  \"output\": {\n    \"properties\": {\n      }\n  }\n}")
			})

			Convey(`pass through user args`, func() {
				// Delimiter inside user args should also be passed through
				expectedUserArgs := []string{"foo", "bar", ArgsDelim, "baz"}
				Convey(`when output is not specified`, func() {
					args = append(args, ArgsDelim)
					args = append(args, expectedUserArgs...)
					exitcode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
						So(userArgs, ShouldResemble, expectedUserArgs)
						return nil
					})
					So(exitcode, ShouldEqual, 0)
				})
				Convey(`when output is specified`, func() {
					tdir, err := ioutil.TempDir("", "luciexe-exe-test")
					So(err, ShouldBeNil)
					defer os.RemoveAll(tdir)
					args = append(args, luciexe.OutputCLIArg, filepath.Join(tdir, "out.pb"), ArgsDelim)
					args = append(args, expectedUserArgs...)
					exitcode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
						So(userArgs, ShouldResemble, expectedUserArgs)
						return nil
					})
					So(exitcode, ShouldEqual, 0)
				})
			})
		})
	})
}
<commit_msg>luciexe: Add test to validate output on error<commit_after>// Copyright 2019 The LUCI 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 exe

import (
	"bytes"
	"compress/zlib"
	"context"
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"

	"github.com/golang/protobuf/proto"
	structpb "github.com/golang/protobuf/ptypes/struct"

	bbpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/logdog/client/butlerlib/bootstrap"
	"go.chromium.org/luci/logdog/client/butlerlib/streamclient"
	"go.chromium.org/luci/luciexe"

	. "github.com/smartystreets/goconvey/convey"
	. "go.chromium.org/luci/common/testing/assertions"
)

func TestExe(t *testing.T) {
	t.Parallel()

	Convey(`test exe`, t, func() {
		client := streamclient.NewFake("test_namespace")
		ldClient := &bootstrap.Bootstrap{
			CoordinatorHost: "test.example.com",
			Project:         "test_project",
			Prefix:          "test_prefix",
			Namespace:       "test_namespace",
			Client:          client.Client,
		}
		var ldErr error
		bootstrapGet := func() (*bootstrap.Bootstrap, error) { return ldClient, ldErr }

		getBuilds := func(decompress bool) []*bbpb.Build {
			fakeData := client.GetFakeData()["test_namespace/build.proto"]
			if decompress {
				So(fakeData.GetFlags().ContentType, ShouldEqual, luciexe.BuildProtoZlibContentType)
			} else {
				So(fakeData.GetFlags().ContentType, ShouldEqual, luciexe.BuildProtoContentType)
			}

			dgs := fakeData.GetDatagrams()
			So(len(dgs), ShouldBeGreaterThanOrEqualTo, 1)

			ret := make([]*bbpb.Build, len(dgs))
			for i, dg := range dgs {
				ret[i] = &bbpb.Build{}

				var data []byte

				if decompress {
					r, err := zlib.NewReader(bytes.NewBufferString(dg))
					So(err, ShouldBeNil)
					data, err = ioutil.ReadAll(r)
					So(err, ShouldBeNil)
				} else {
					data = []byte(dg)
				}

				So(proto.Unmarshal(data, ret[i]), ShouldBeNil)
			}
			return ret
		}
		lastBuild := func() *bbpb.Build {
			builds := getBuilds(false)
			return builds[len(builds)-1]
		}

		args := []string{"fake_test_executable"}

		ctx := context.Background()

		Convey(`basic`, func() {
			Convey(`success`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.Status = bbpb.Status_SCHEDULED
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status: bbpb.Status_SUCCESS,
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`failure`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					return errors.New("bad stuff")
				})
				So(exitCode, ShouldEqual, 1)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_FAILURE,
					SummaryMarkdown: "Final error: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`infra failure`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					return errors.New("bad stuff", InfraErrorTag)
				})
				So(exitCode, ShouldEqual, 1)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "Final infra error: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`panic`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					panic(errors.New("bad stuff"))
				})
				So(exitCode, ShouldEqual, 2)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "Final panic: bad stuff",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})

			Convey(`respect user program status`, func() {
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.Status = bbpb.Status_INFRA_FAILURE
					build.SummaryMarkdown = "status set inside"
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_INFRA_FAILURE,
					SummaryMarkdown: "status set inside",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
			})
		})

		Convey(`send`, func() {
			exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
				build.SummaryMarkdown = "Hi. I did stuff."
				bs()
				return errors.New("oh no i failed")
			})
			So(exitCode, ShouldEqual, 1)
			builds := getBuilds(false)
			So(len(builds), ShouldEqual, 2)
			So(builds[0], ShouldResembleProto, &bbpb.Build{
				SummaryMarkdown: "Hi. I did stuff.",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
			So(builds[len(builds)-1], ShouldResembleProto, &bbpb.Build{
				Status:          bbpb.Status_FAILURE,
				SummaryMarkdown: "Hi. I did stuff.\n\nFinal error: oh no i failed",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
		})

		Convey(`send (zlib)`, func() {
			exitCode := runCtx(ctx, args, bootstrapGet, []Option{WithZlibCompression(5)}, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
				build.SummaryMarkdown = "Hi. I did stuff."
				bs()
				return errors.New("oh no i failed")
			})
			So(exitCode, ShouldEqual, 1)
			builds := getBuilds(true)
			So(len(builds), ShouldEqual, 2)
			So(builds[0], ShouldResembleProto, &bbpb.Build{
				SummaryMarkdown: "Hi. I did stuff.",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
			So(builds[len(builds)-1], ShouldResembleProto, &bbpb.Build{
				Status:          bbpb.Status_FAILURE,
				SummaryMarkdown: "Hi. I did stuff.\n\nFinal error: oh no i failed",
				Output: &bbpb.Build_Output{
					Properties: &structpb.Struct{},
				},
			})
		})

		Convey(`output`, func() {
			tdir, err := ioutil.TempDir("", "luciexe-exe-test")
			So(err, ShouldBeNil)
			defer os.RemoveAll(tdir)

			Convey(`binary`, func() {
				outFile := filepath.Join(tdir, "out.pb")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					err := WriteProperties(build.Output.Properties, map[string]interface{}{
						"some": "thing",
					})
					if err != nil {
						panic(err)
					}

					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{
							Fields: map[string]*structpb.Value{
								"some": {Kind: &structpb.Value_StringValue{
									StringValue: "thing",
								}},
							},
						},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"`\f\x82\x01\x13\n\x11\n\x0f\n\x04some\x12\a\x1a\x05thing\xa2\x01\x03Hi.")
			})

			Convey(`textpb`, func() {
				outFile := filepath.Join(tdir, "out.textpb")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"status: SUCCESS\nsummary_markdown: \"Hi.\"\noutput: <\n  properties: <\n  >\n>\n")
			})

			Convey(`jsonpb`, func() {
				outFile := filepath.Join(tdir, "out.json")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					return nil
				})
				So(exitCode, ShouldEqual, 0)
				So(lastBuild(), ShouldResembleProto, &bbpb.Build{
					Status:          bbpb.Status_SUCCESS,
					SummaryMarkdown: "Hi.",
					Output: &bbpb.Build_Output{
						Properties: &structpb.Struct{},
					},
				})
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"{\n  \"status\": \"SUCCESS\",\n  \"summary_markdown\": \"Hi.\",\n  \"output\": {\n    \"properties\": {\n      }\n  }\n}")
			})

			Convey(`pass through user args`, func() {
				// Delimiter inside user args should also be passed through
				expectedUserArgs := []string{"foo", "bar", ArgsDelim, "baz"}
				Convey(`when output is not specified`, func() {
					args = append(args, ArgsDelim)
					args = append(args, expectedUserArgs...)
					exitcode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
						So(userArgs, ShouldResemble, expectedUserArgs)
						return nil
					})
					So(exitcode, ShouldEqual, 0)
				})
				Convey(`when output is specified`, func() {
					tdir, err := ioutil.TempDir("", "luciexe-exe-test")
					So(err, ShouldBeNil)
					defer os.RemoveAll(tdir)
					args = append(args, luciexe.OutputCLIArg, filepath.Join(tdir, "out.pb"), ArgsDelim)
					args = append(args, expectedUserArgs...)
					exitcode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
						So(userArgs, ShouldResemble, expectedUserArgs)
						return nil
					})
					So(exitcode, ShouldEqual, 0)
				})
			})

			Convey(`write output on error`, func() {
				outFile := filepath.Join(tdir, "out.json")
				args = append(args, luciexe.OutputCLIArg, outFile)
				exitCode := runCtx(ctx, args, bootstrapGet, nil, func(ctx context.Context, build *bbpb.Build, userArgs []string, bs BuildSender) error {
					build.SummaryMarkdown = "Hi."
					return errors.New("bad stuff")
				})
				So(exitCode, ShouldEqual, 1)
				data, err := ioutil.ReadFile(outFile)
				So(err, ShouldBeNil)
				So(string(data), ShouldResemble,
					"{\n  \"status\": \"FAILURE\",\n  \"summary_markdown\": \"Hi.\\n\\nFinal error: bad stuff\",\n  \"output\": {\n    \"properties\": {\n      }\n  }\n}")
			})
		})
	})
}
<|endoftext|>


<commit_before><commit_msg>Skip bind route service test under client credentials mode<commit_after><|endoftext|>


<commit_before>package ploop

// A test suite, also serving as an example of how to use the package

import (
	"github.com/dustin/go-humanize"
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

var (
	old_pwd  string
	test_dir string
	d        Ploop
	snap     string
)

const baseDelta = "root.hdd"

func TestPrepare(t *testing.T) {
	var err error

	old_pwd, err = os.Getwd()
	if err != nil {
		t.Fatalf("Getwd: %s", err)
	}

	test_dir, err = ioutil.TempDir(old_pwd, "tmp-test")
	if err != nil {
		t.Fatalf("TempDir %q: %s", old_pwd, err)
	}

	err = os.Chdir(test_dir)
	if err != nil {
		t.Fatalf("Chdir %q: %s", test_dir, err)
	}

	SetVerboseLevel(255)
	SetLogLevel(1)
	SetLogFile("ploop.log")

	t.Logf("Running tests in %s", test_dir)
}

func TestUUID(t *testing.T) {
	uuid, e := UUID()
	if e != nil {
		t.Errorf("UUID: %s", e)
	}

	t.Logf("Got uuid %s", uuid)
}

func TestCreate(t *testing.T) {
	size := "384M"
	var p CreateParam

	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}
	p.Size = s
	p.File = baseDelta

	e = Create(&p)
	if e != nil {
		t.Fatalf("Create: %s", e)
	}
}

func TestOpen(t *testing.T) {
	var e error

	d, e = Open("DiskDescriptor.xml")
	if e != nil {
		t.Errorf("Open: %s ", e)
	}
}

func TestMount(t *testing.T) {
	mnt := "mnt"

	e := os.Mkdir(mnt, 0755)
	if e != nil {
		t.Fatalf("os.Mkdir: %s", e)
	}

	p := MountParam{Target: mnt}
	dev, e := d.Mount(&p)
	if e != nil {
		t.Fatalf("Open: %s", e)
	} else {
		t.Logf("Mounted; ploop device %s", dev)
	}
}

func resize(t *testing.T, size string, offline bool) {
	if offline && testing.Short() {
		t.Skip("skipping offline resize test in short mode.")
	}
	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}

	e = d.Resize(s, offline)
	if e != nil {
		t.Fatalf("Resize to %s (%d) failed: %s", size, s, e)
	}
}

func TestResizeOnlineShrink(t *testing.T) {
	resize(t, "256MB", false)
}

func TestResizeOnlineGrow(t *testing.T) {
	resize(t, "512MB", false)
}

func TestSnapshot(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created online Snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func copyFile(src, dst string) error {

	cmd := exec.Command("cp", "-a", src, dst)
	err := cmd.Run()

	return err
}

func testReplace(t *testing.T) {
	var p ReplaceParam
	newDelta := baseDelta + ".new"
	e := copyFile(baseDelta, newDelta)
	if e != nil {
		t.Fatalf("copyFile: %s", e)
	}

	p.File = newDelta
	p.CurFile = baseDelta
	p.Flags = KeepName
	e = d.Replace(&p)
	if e != nil {
		t.Fatalf("Replace: %s", e)
	}
}

func TestReplaceOnline(t *testing.T) {
	testReplace(t)
}

func TestSnapshotDelete(t *testing.T) {
	e := d.DeleteSnapshot(snap)
	if e != nil {
		t.Fatalf("DeleteSnapshot: %s", e)
	} else {
		t.Logf("Deleted snapshot %s", snap)
	}
}

func TestUmount(t *testing.T) {
	e := d.Umount()
	if e != nil {
		t.Fatalf("Umount: %s", e)
	}
}

func TestResizeOfflineShrink(t *testing.T) {
	resize(t, "256MB", true)
}

func TestResizeOfflineGrow(t *testing.T) {
	resize(t, "512MB", true)
}

func TestResizeOfflineShrinkAgain(t *testing.T) {
	resize(t, "256MB", true)
}

func TestSnapshotOffline(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created offline Snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func TestReplaceOffline(t *testing.T) {
	testReplace(t)
}

func TestSnapshotSwitch(t *testing.T) {
	e := d.SwitchSnapshot(snap)
	if e != nil {
		t.Fatalf("SwitchSnapshot: %s", e)
	} else {
		t.Logf("Switched to snapshot %s", snap)
	}
}

func TestFSInfo(t *testing.T) {
	i, e := FSInfo("DiskDescriptor.xml")

	if e != nil {
		t.Errorf("FSInfo: %v", e)
	} else {
		bTotal := i.Blocks * i.Blocksize
		bAvail := i.Blocks_free * i.Blocksize
		bUsed := bTotal - bAvail

		iTotal := i.Inodes
		iAvail := i.Inodes_free
		iUsed := iTotal - iAvail

		t.Logf("\n             Size       Used      Avail Use%%\n%7s %9s %10s %10s %3d%%\n%7s %9d %10d %10d %3d%%",
			"Blocks",
			humanize.Bytes(bTotal),
			humanize.Bytes(bUsed),
			humanize.Bytes(bAvail),
			100*bUsed/bTotal,
			"Inodes",
			iTotal,
			iUsed,
			iAvail,
			100*iUsed/iTotal)
		t.Logf("\nInode ratio: 1 inode per %s of disk space",
			humanize.Bytes(bTotal/iTotal))
	}
}

func TestImageInfo(t *testing.T) {
	i, e := d.ImageInfo()
	if e != nil {
		t.Errorf("ImageInfo: %v", e)
	} else {
		t.Logf("\n              Blocks  Blocksize       Size  Ver\n%20d %10d %10s %4d",
			i.Blocks, i.Blocksize,
			humanize.Bytes(512*i.Blocks),
			i.Version)
	}

}

// TestCleanup is the last test, removing files created by previous tests
func TestCleanup(t *testing.T) {
	d.Close()
	os.Chdir(old_pwd)
	os.RemoveAll(test_dir)
}
<commit_msg>copyFile: simplify<commit_after>package ploop

// A test suite, also serving as an example of how to use the package

import (
	"github.com/dustin/go-humanize"
	"io/ioutil"
	"os"
	"os/exec"
	"testing"
)

var (
	old_pwd  string
	test_dir string
	d        Ploop
	snap     string
)

const baseDelta = "root.hdd"

func TestPrepare(t *testing.T) {
	var err error

	old_pwd, err = os.Getwd()
	if err != nil {
		t.Fatalf("Getwd: %s", err)
	}

	test_dir, err = ioutil.TempDir(old_pwd, "tmp-test")
	if err != nil {
		t.Fatalf("TempDir %q: %s", old_pwd, err)
	}

	err = os.Chdir(test_dir)
	if err != nil {
		t.Fatalf("Chdir %q: %s", test_dir, err)
	}

	SetVerboseLevel(255)
	SetLogLevel(1)
	SetLogFile("ploop.log")

	t.Logf("Running tests in %s", test_dir)
}

func TestUUID(t *testing.T) {
	uuid, e := UUID()
	if e != nil {
		t.Errorf("UUID: %s", e)
	}

	t.Logf("Got uuid %s", uuid)
}

func TestCreate(t *testing.T) {
	size := "384M"
	var p CreateParam

	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}
	p.Size = s
	p.File = baseDelta

	e = Create(&p)
	if e != nil {
		t.Fatalf("Create: %s", e)
	}
}

func TestOpen(t *testing.T) {
	var e error

	d, e = Open("DiskDescriptor.xml")
	if e != nil {
		t.Errorf("Open: %s ", e)
	}
}

func TestMount(t *testing.T) {
	mnt := "mnt"

	e := os.Mkdir(mnt, 0755)
	if e != nil {
		t.Fatalf("os.Mkdir: %s", e)
	}

	p := MountParam{Target: mnt}
	dev, e := d.Mount(&p)
	if e != nil {
		t.Fatalf("Open: %s", e)
	} else {
		t.Logf("Mounted; ploop device %s", dev)
	}
}

func resize(t *testing.T, size string, offline bool) {
	if offline && testing.Short() {
		t.Skip("skipping offline resize test in short mode.")
	}
	s, e := humanize.ParseBytes(size)
	if e != nil {
		t.Fatalf("humanize.ParseBytes: can't parse %s: %s", size, e)
	}

	e = d.Resize(s, offline)
	if e != nil {
		t.Fatalf("Resize to %s (%d) failed: %s", size, s, e)
	}
}

func TestResizeOnlineShrink(t *testing.T) {
	resize(t, "256MB", false)
}

func TestResizeOnlineGrow(t *testing.T) {
	resize(t, "512MB", false)
}

func TestSnapshot(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created online Snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func copyFile(src, dst string) error {
	return exec.Command("cp", "-a", src, dst).Run
}

func testReplace(t *testing.T) {
	var p ReplaceParam
	newDelta := baseDelta + ".new"
	e := copyFile(baseDelta, newDelta)
	if e != nil {
		t.Fatalf("copyFile: %s", e)
	}

	p.File = newDelta
	p.CurFile = baseDelta
	p.Flags = KeepName
	e = d.Replace(&p)
	if e != nil {
		t.Fatalf("Replace: %s", e)
	}
}

func TestReplaceOnline(t *testing.T) {
	testReplace(t)
}

func TestSnapshotDelete(t *testing.T) {
	e := d.DeleteSnapshot(snap)
	if e != nil {
		t.Fatalf("DeleteSnapshot: %s", e)
	} else {
		t.Logf("Deleted snapshot %s", snap)
	}
}

func TestUmount(t *testing.T) {
	e := d.Umount()
	if e != nil {
		t.Fatalf("Umount: %s", e)
	}
}

func TestResizeOfflineShrink(t *testing.T) {
	resize(t, "256MB", true)
}

func TestResizeOfflineGrow(t *testing.T) {
	resize(t, "512MB", true)
}

func TestResizeOfflineShrinkAgain(t *testing.T) {
	resize(t, "256MB", true)
}

func TestSnapshotOffline(t *testing.T) {
	uuid, e := d.Snapshot()
	if e != nil {
		t.Fatalf("Snapshot: %s", e)
	} else {
		t.Logf("Created offline Snapshot; uuid %s", uuid)
	}

	snap = uuid
}

func TestReplaceOffline(t *testing.T) {
	testReplace(t)
}

func TestSnapshotSwitch(t *testing.T) {
	e := d.SwitchSnapshot(snap)
	if e != nil {
		t.Fatalf("SwitchSnapshot: %s", e)
	} else {
		t.Logf("Switched to snapshot %s", snap)
	}
}

func TestFSInfo(t *testing.T) {
	i, e := FSInfo("DiskDescriptor.xml")

	if e != nil {
		t.Errorf("FSInfo: %v", e)
	} else {
		bTotal := i.Blocks * i.Blocksize
		bAvail := i.Blocks_free * i.Blocksize
		bUsed := bTotal - bAvail

		iTotal := i.Inodes
		iAvail := i.Inodes_free
		iUsed := iTotal - iAvail

		t.Logf("\n             Size       Used      Avail Use%%\n%7s %9s %10s %10s %3d%%\n%7s %9d %10d %10d %3d%%",
			"Blocks",
			humanize.Bytes(bTotal),
			humanize.Bytes(bUsed),
			humanize.Bytes(bAvail),
			100*bUsed/bTotal,
			"Inodes",
			iTotal,
			iUsed,
			iAvail,
			100*iUsed/iTotal)
		t.Logf("\nInode ratio: 1 inode per %s of disk space",
			humanize.Bytes(bTotal/iTotal))
	}
}

func TestImageInfo(t *testing.T) {
	i, e := d.ImageInfo()
	if e != nil {
		t.Errorf("ImageInfo: %v", e)
	} else {
		t.Logf("\n              Blocks  Blocksize       Size  Ver\n%20d %10d %10s %4d",
			i.Blocks, i.Blocksize,
			humanize.Bytes(512*i.Blocks),
			i.Version)
	}

}

// TestCleanup is the last test, removing files created by previous tests
func TestCleanup(t *testing.T) {
	d.Close()
	os.Chdir(old_pwd)
	os.RemoveAll(test_dir)
}
<|endoftext|>


<commit_before>package registry

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/codeskyblue/go-sh"
	"github.com/dokku/dokku/plugins/common"
)

func getRegistryServerForApp(appName string) string {
	value := common.PropertyGet("registry", appName, "server")
	if value == "" {
		value = common.PropertyGet("registry", "--global", "server")
	}

	value = strings.TrimSuffix(value, "/") + "/"
	if value == "hub.docker.com/" || value == "docker.io/" {
		value = ""
	}

	return value
}

func isPushEnabled(appName string) bool {
	return reportComputedPushOnRelease(appName) == "true"
}

func incrementTagVersion(appName string) (string, error) {
	tag := common.PropertyGet("registry", appName, "tag-version")
	if tag == "" {
		tag = "0"
	}

	version, err := strconv.Atoi(tag)
	if err != nil {
		return "", fmt.Errorf("Unable to convert existing tag version (%s) to integer: %v", tag, err)
	}

	version++
	common.LogVerboseQuiet(fmt.Sprintf("Bumping tag to %d", version))
	if err = common.PropertyWrite("registry", appName, "tag-version", strconv.Itoa(version)); err != nil {
		return "", err
	}

	return strconv.Itoa(version), nil
}

func pushToRegistry(appName string, tag string) error {
	common.LogVerboseQuiet("Retrieving image info for app")

	registryServer := getRegistryServerForApp(appName)
	imageRepo := reportComputedImageRepo(appName)
	imageTag, _ := common.GetRunningImageTag(appName)
	image := common.GetAppImageName(appName, imageTag, "")
	imageID, _ := common.DockerInspect(image, "{{ .Id }}")

	common.LogVerboseQuiet(fmt.Sprintf("Tagging $IMAGE_REPO:%s in registry format", tag))
	if !dockerTag(imageID, fmt.Sprintf("%s%s:%s", registryServer, imageRepo, tag)) {
		// TODO: better error
		return errors.New("Unable to tag image")
	}

	if !dockerTag(imageID, fmt.Sprintf("%s:%s", imageRepo, tag)) {
		// TODO: better error
		return errors.New("Unable to tag image")
	}

	// fn-registry-create-repository "$APP" "$DOKKU_REGISTRY_SERVER" "$IMAGE_REPO"

	common.LogVerboseQuiet("Pushing $IMAGE_REPO:$TAG")
	// docker push "${DOKKU_REGISTRY_SERVER}${IMAGE_REPO}:${TAG}"

	common.LogVerboseQuiet("Cleaning up")
	// fn-registry-image-cleanup "$APP" "$DOKKU_REGISTRY_SERVER" "$IMAGE_REPO" "$IMAGE_TAG" "$TAG"

	common.LogVerboseQuiet("Image $IMAGE_REPO:$TAG pushed")
	return nil
}

func dockerTag(imageID string, imageTag string) bool {
	cmd := sh.Command(common.DockerBin(), "image", "tag", imageID, imageTag)
	cmd.Stdout = nil
	cmd.Stderr = nil
	if err := cmd.Run(); err != nil {
		return false
	}

	return true
}
<commit_msg>feat: push tag<commit_after>package registry

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/codeskyblue/go-sh"
	"github.com/dokku/dokku/plugins/common"
)

func getRegistryServerForApp(appName string) string {
	value := common.PropertyGet("registry", appName, "server")
	if value == "" {
		value = common.PropertyGet("registry", "--global", "server")
	}

	value = strings.TrimSuffix(value, "/") + "/"
	if value == "hub.docker.com/" || value == "docker.io/" {
		value = ""
	}

	return value
}

func isPushEnabled(appName string) bool {
	return reportComputedPushOnRelease(appName) == "true"
}

func incrementTagVersion(appName string) (string, error) {
	tag := common.PropertyGet("registry", appName, "tag-version")
	if tag == "" {
		tag = "0"
	}

	version, err := strconv.Atoi(tag)
	if err != nil {
		return "", fmt.Errorf("Unable to convert existing tag version (%s) to integer: %v", tag, err)
	}

	version++
	common.LogVerboseQuiet(fmt.Sprintf("Bumping tag to %d", version))
	if err = common.PropertyWrite("registry", appName, "tag-version", strconv.Itoa(version)); err != nil {
		return "", err
	}

	return strconv.Itoa(version), nil
}

func pushToRegistry(appName string, tag string) error {
	common.LogVerboseQuiet("Retrieving image info for app")

	registryServer := getRegistryServerForApp(appName)
	imageRepo := reportComputedImageRepo(appName)
	imageTag, _ := common.GetRunningImageTag(appName)
	image := common.GetAppImageName(appName, imageTag, "")
	imageID, _ := common.DockerInspect(image, "{{ .Id }}")

	common.LogVerboseQuiet(fmt.Sprintf("Tagging $IMAGE_REPO:%s in registry format", tag))
	if !dockerTag(imageID, fmt.Sprintf("%s%s:%s", registryServer, imageRepo, tag)) {
		// TODO: better error
		return errors.New("Unable to tag image")
	}

	if !dockerTag(imageID, fmt.Sprintf("%s:%s", imageRepo, tag)) {
		// TODO: better error
		return errors.New("Unable to tag image")
	}

	// fn-registry-create-repository "$APP" "$DOKKU_REGISTRY_SERVER" "$IMAGE_REPO"

	common.LogVerboseQuiet("Pushing $IMAGE_REPO:$TAG")
	if !dockerPush(fmt.Sprintf("%s%s:%s", registryServer, imageRepo, tag)) {
		// TODO: better error
		return errors.New("Unable to push image")
	}

	common.LogVerboseQuiet("Cleaning up")
	// fn-registry-image-cleanup "$APP" "$DOKKU_REGISTRY_SERVER" "$IMAGE_REPO" "$IMAGE_TAG" "$TAG"

	common.LogVerboseQuiet("Image $IMAGE_REPO:$TAG pushed")
	return nil
}

func dockerTag(imageID string, imageTag string) bool {
	cmd := sh.Command(common.DockerBin(), "image", "tag", imageID, imageTag)
	cmd.Stdout = nil
	cmd.Stderr = nil
	if err := cmd.Run(); err != nil {
		return false
	}

	return true
}

func dockerPush(imageTag string) bool {
	cmd := sh.Command(common.DockerBin(), "image", "push", imageTag)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return false
	}

	return true
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iron-io/golog"
)

type GlockConfig struct {
	Port      int `json:"port"`
	Logging   LoggingConfig
	LockLimit int64 `json:"lock_limit"`
}

type LoggingConfig struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}

type timeoutLock struct {
	mutex     sync.Mutex
	id        int64 // unique ID of the current lock. Only allow an unlock if the correct id is passed
	lockCount int64
}

var locksLock sync.RWMutex
var locks = map[string]*timeoutLock{}
var config GlockConfig

func main() {
	var port int
	var configFile string
	flag.IntVar(&port, "p", 45625, "port")
	flag.StringVar(&configFile, "c", "", "Name of the the file that contains config information")
	flag.Parse()

	if configFile != "" {
		LoadConfig(configFile, &config)
	}

	if config.Port == 0 {
		config.Port = port
	}

	if config.Logging.Level == "" {
		config.Logging.Level = "debug"
	}

	listener, err := net.Listen("tcp", ":"+strconv.Itoa(config.Port))
	if err != nil {
		log.Fatalln("error listening", err)
	}

	golog.SetLogLevel(config.Logging.Level)
	golog.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	golog.Infoln("Glock Server available at port ", config.Port)

	for {
		conn, err := listener.Accept()
		if err != nil {
			golog.Errorln("error accepting", err)
			return
		}
		go handleConn(conn)
	}
}

var (
	unlockedResponse    = []byte("UNLOCKED\r\n")
	notUnlockedResponse = []byte("NOT_UNLOCKED\r\n")
	pongResponse        = []byte("PONG\r\n")

	errBadFormat      = []byte("ERROR 400 bad command format\r\n")
	errUnknownCommand = []byte("ERROR 405 unknown command\r\n")
	errLockNotFound   = []byte("ERROR 404 lock not found\r\n")
	errLockAtCapacity = []byte("ERROR 503 lock at capacity\r\n")
)

func handleConn(conn net.Conn) {
	defer func() {
		conn.Close()
		// make sure a panic doesn't take down the whole server
		err := recover()
		if err != nil {
			golog.Errorf("recovered from panic: %v\n%s\n", err, debug.Stack())
		}
	}()

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		split := strings.Fields(scanner.Text())

		if split[0] == "PING" {
			conn.Write(pongResponse)
			golog.Debugf("PING PONG")
			continue
		}

		if len(split) < 3 {
			conn.Write(errBadFormat)
			continue
		}

		cmd := split[0]
		key := split[1]
		switch cmd {
		// LOCK <key> <timeout>
		case "LOCK":
			timeout, err := strconv.Atoi(split[2])

			if err != nil {
				conn.Write(errBadFormat)
				golog.Errorln(string(errBadFormat), ": ", split)
				continue
			}
			locksLock.RLock()
			lock, ok := locks[key]
			locksLock.RUnlock()
			if !ok {
				// lock doesn't exist; create it
				locksLock.Lock()
				lock, ok = locks[key]
				if !ok {
					lock = &timeoutLock{}
					locks[key] = lock
				}
				locksLock.Unlock()
			}

			if !lock.lockMutex() {
				conn.Write(errLockAtCapacity)
				continue
			}
			id := atomic.AddInt64(&lock.id, 1)
			time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
				if atomic.CompareAndSwapInt64(&lock.id, id, id+1) {
					lock.unlockMutex()
					golog.Debugf("P %-5d | Timedout: %-12d | Key:  %-15s | Id: %d", config.Port, timeout, key, id)
				}
			})
			fmt.Fprintf(conn, "LOCKED %v\n", id)

			golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Timeout: %dms", config.Port, cmd, key, timeout)
			golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "LOCKED", key, id)

		// UNLOCK <key> <id>
		case "UNLOCK":
			id, err := strconv.ParseInt(split[2], 10, 64)

			if err != nil {
				conn.Write(errBadFormat)
				golog.Errorln(string(errBadFormat), ": ", split)
				continue
			}
			locksLock.RLock()
			lock, ok := locks[key]
			locksLock.RUnlock()
			if !ok {
				conn.Write(errLockNotFound)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Errorln(string(errLockNotFound), ": ", split, "| P ", config.Port)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s", config.Port, "404", key)
				continue
			}
			if atomic.CompareAndSwapInt64(&lock.id, id, id+1) {
				lock.unlockMutex()
				conn.Write(unlockedResponse)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "UNLOCKED", key, id)
			} else {
				conn.Write(notUnlockedResponse)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "NOT_UNLOCKED", key, id)
			}

		default:
			conn.Write(errUnknownCommand)
			golog.Errorln(string(errUnknownCommand), ": ", split)
			continue
		}
	}
}

func LoadConfig(configFile string, config interface{}) {
	config_s, err := ioutil.ReadFile(configFile)
	if err != nil {
		log.Fatalln("Couldn't find config at:", configFile)
	}

	err = json.Unmarshal(config_s, &config)
	if err != nil {
		log.Fatalln("Couldn't unmarshal config!", err)
	}
	golog.Infoln("config:", config)
}

func (l *timeoutLock) lockMutex() bool {
	if config.LockLimit != 0 {
		for {
			count := atomic.LoadInt64(&l.lockCount)
			if count >= config.LockLimit {
				return false
			}

			if atomic.CompareAndSwapInt64(&l.lockCount, count, count+1) {
				break
			}
		}
	}
	l.mutex.Lock()
	return true
}

func (l *timeoutLock) unlockMutex() {
	l.mutex.Unlock()
	if config.LockLimit != 0 {
		atomic.AddInt64(&l.lockCount, -1)
	}
}
<commit_msg>gofmt<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"runtime/debug"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iron-io/golog"
)

type GlockConfig struct {
	Port      int   `json:"port"`
	LockLimit int64 `json:"lock_limit"`
	Logging   LoggingConfig
}

type LoggingConfig struct {
	To     string `json:"to"`
	Level  string `json:"level"`
	Prefix string `json:"prefix"`
}

type timeoutLock struct {
	mutex     sync.Mutex
	id        int64 // unique ID of the current lock. Only allow an unlock if the correct id is passed
	lockCount int64
}

var locksLock sync.RWMutex
var locks = map[string]*timeoutLock{}
var config GlockConfig

func main() {
	var port int
	var configFile string
	flag.IntVar(&port, "p", 45625, "port")
	flag.StringVar(&configFile, "c", "", "Name of the the file that contains config information")
	flag.Parse()

	if configFile != "" {
		LoadConfig(configFile, &config)
	}

	if config.Port == 0 {
		config.Port = port
	}

	if config.Logging.Level == "" {
		config.Logging.Level = "debug"
	}

	listener, err := net.Listen("tcp", ":"+strconv.Itoa(config.Port))
	if err != nil {
		log.Fatalln("error listening", err)
	}

	golog.SetLogLevel(config.Logging.Level)
	golog.SetLogLocation(config.Logging.To, config.Logging.Prefix)

	golog.Infoln("Glock Server available at port ", config.Port)

	for {
		conn, err := listener.Accept()
		if err != nil {
			golog.Errorln("error accepting", err)
			return
		}
		go handleConn(conn)
	}
}

var (
	unlockedResponse    = []byte("UNLOCKED\r\n")
	notUnlockedResponse = []byte("NOT_UNLOCKED\r\n")
	pongResponse        = []byte("PONG\r\n")

	errBadFormat      = []byte("ERROR 400 bad command format\r\n")
	errUnknownCommand = []byte("ERROR 405 unknown command\r\n")
	errLockNotFound   = []byte("ERROR 404 lock not found\r\n")
	errLockAtCapacity = []byte("ERROR 503 lock at capacity\r\n")
)

func handleConn(conn net.Conn) {
	defer func() {
		conn.Close()
		// make sure a panic doesn't take down the whole server
		err := recover()
		if err != nil {
			golog.Errorf("recovered from panic: %v\n%s\n", err, debug.Stack())
		}
	}()

	scanner := bufio.NewScanner(conn)
	for scanner.Scan() {
		split := strings.Fields(scanner.Text())

		if split[0] == "PING" {
			conn.Write(pongResponse)
			golog.Debugf("PING PONG")
			continue
		}

		if len(split) < 3 {
			conn.Write(errBadFormat)
			continue
		}

		cmd := split[0]
		key := split[1]
		switch cmd {
		// LOCK <key> <timeout>
		case "LOCK":
			timeout, err := strconv.Atoi(split[2])

			if err != nil {
				conn.Write(errBadFormat)
				golog.Errorln(string(errBadFormat), ": ", split)
				continue
			}
			locksLock.RLock()
			lock, ok := locks[key]
			locksLock.RUnlock()
			if !ok {
				// lock doesn't exist; create it
				locksLock.Lock()
				lock, ok = locks[key]
				if !ok {
					lock = &timeoutLock{}
					locks[key] = lock
				}
				locksLock.Unlock()
			}

			if !lock.lockMutex() {
				conn.Write(errLockAtCapacity)
				continue
			}
			id := atomic.AddInt64(&lock.id, 1)
			time.AfterFunc(time.Duration(timeout)*time.Millisecond, func() {
				if atomic.CompareAndSwapInt64(&lock.id, id, id+1) {
					lock.unlockMutex()
					golog.Debugf("P %-5d | Timedout: %-12d | Key:  %-15s | Id: %d", config.Port, timeout, key, id)
				}
			})
			fmt.Fprintf(conn, "LOCKED %v\n", id)

			golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Timeout: %dms", config.Port, cmd, key, timeout)
			golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "LOCKED", key, id)

		// UNLOCK <key> <id>
		case "UNLOCK":
			id, err := strconv.ParseInt(split[2], 10, 64)

			if err != nil {
				conn.Write(errBadFormat)
				golog.Errorln(string(errBadFormat), ": ", split)
				continue
			}
			locksLock.RLock()
			lock, ok := locks[key]
			locksLock.RUnlock()
			if !ok {
				conn.Write(errLockNotFound)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Errorln(string(errLockNotFound), ": ", split, "| P ", config.Port)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s", config.Port, "404", key)
				continue
			}
			if atomic.CompareAndSwapInt64(&lock.id, id, id+1) {
				lock.unlockMutex()
				conn.Write(unlockedResponse)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "UNLOCKED", key, id)
			} else {
				conn.Write(notUnlockedResponse)

				golog.Debugf("P %-5d | Request:  %-12s | Key:  %-15s | Id: %d", config.Port, cmd, key, id)
				golog.Debugf("P %-5d | Response: %-12s | Key:  %-15s | Id: %d", config.Port, "NOT_UNLOCKED", key, id)
			}

		default:
			conn.Write(errUnknownCommand)
			golog.Errorln(string(errUnknownCommand), ": ", split)
			continue
		}
	}
}

func LoadConfig(configFile string, config interface{}) {
	config_s, err := ioutil.ReadFile(configFile)
	if err != nil {
		log.Fatalln("Couldn't find config at:", configFile)
	}

	err = json.Unmarshal(config_s, &config)
	if err != nil {
		log.Fatalln("Couldn't unmarshal config!", err)
	}
	golog.Infoln("config:", config)
}

func (l *timeoutLock) lockMutex() bool {
	if config.LockLimit != 0 {
		for {
			count := atomic.LoadInt64(&l.lockCount)
			if count >= config.LockLimit {
				return false
			}

			if atomic.CompareAndSwapInt64(&l.lockCount, count, count+1) {
				break
			}
		}
	}
	l.mutex.Lock()
	return true
}

func (l *timeoutLock) unlockMutex() {
	l.mutex.Unlock()
	if config.LockLimit != 0 {
		atomic.AddInt64(&l.lockCount, -1)
	}
}
<|endoftext|>


<commit_before>package QesyDb

import(
    "database/sql"
    "fmt"
    "strings"
    "log"
    _"github.com/go-sql-driver/mysql"
    "reflect"
    "strconv"
    "time"
)

var Db *sql.DB

type Model struct{
    Cond interface{}
    Insert map[string]string
    Update map[string]string
    Field string
    Table string
    Index string
    Limit interface{}
    Sort string
    Fetch int
}

// Connect  is a method with a sql.
func Connect(connStr string){
    sqlDb, err := sql.Open("mysql", connStr)
    if(err != nil){
        log.Fatal("mysql connect error")
    }
    err = sqlDb.Ping()
    if(err != nil){
        log.Fatal("mysql ping error")
    }
	fmt.Println("mysql connect sueccss")
    Db = sqlDb
}

// ExecSelectIndex  is a method with a sql.
func (m *Model) ExecSelectIndex() (map[string]map[string]string, error) {
    resultsSlice, _ := m.ExecSelect()
    retArr := map[string]map[string]string{}
    for _, v := range resultsSlice{
        if v[m.Index] == "" {
            continue
        }
        retArr[v[m.Index]] = v
    }
    m.Clean()
    return retArr, nil
}

// ExecSelect is a method with a sql.
func (m *Model) ExecSelect() ([]map[string]string, error) {
    cond := m.getSQLCond()
    field := m.getSQLField()
    sort := m.getSort()
    limit := m.getSQLLimite()
    sql := "SELECT "+ field +" FROM "+m.Table+cond+sort+limit+";"
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    rows, err := stmt.Query()
    if(err != nil){
        return nil, err
    }
    defer rows.Close()
    fields, err := rows.Columns()
    if(err != nil){
        return nil, err
    }
    var resultsSlice []map[string]string
    for rows.Next(){
        result := make(map[string]string)
        var scanResultContainers []interface{}
        for i := 0; i < len(fields); i++ {
            var scanResultContainer interface{}
			scanResultContainers = append(scanResultContainers, &scanResultContainer)
        }
        if err := rows.Scan(scanResultContainers...); err != nil {
			return nil, err
		}
        for k, v := range fields {
            rawValue := reflect.Indirect(reflect.ValueOf(scanResultContainers[k]))
            if rawValue.Interface() == nil{
                continue
            }
            rawType := reflect.TypeOf(rawValue.Interface())
			rawVal := reflect.ValueOf(rawValue.Interface())
            var str string
            switch rawType.Kind() {
                case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
                    str = strconv.FormatInt(rawVal.Int(), 10)
                    result[v] = str
                case reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
                    str = strconv.FormatUint(rawVal.Uint(), 10)
                    result[v] = str
                case reflect.Float32, reflect.Float64:
                    str = strconv.FormatFloat(rawVal.Float(), 'f', -1, 64)
                    result[v] = str
                case reflect.Slice:
                    if rawType.Elem().Kind() == reflect.Uint8 {
                        result[v] = string(rawVal.Interface().([]byte))
                        break
                    }
                case reflect.String:
                    str = rawVal.String()
                    result[v] = str
                case reflect.Struct:
                    str = rawVal.Interface().(time.Time).Format("2006-01-02 15:04:05.000 -0700")
                    result[v] = str
                case reflect.Bool:
                    if rawVal.Bool() {
                        result[v] = "1"
                    } else {
                        result[v] = "0"
                    }
                }
            }
            resultsSlice = append(resultsSlice, result)
        }
        m.Clean()
        return resultsSlice, nil
}

func (m *Model) ExecSelectOne() (map[string]string, error) {
    resultsSlice, _ := m.ExecSelect()
    m.Clean()
    return resultsSlice[0], nil
}

func (m *Model) ExecUpdate() (sql.Result, error){
    updateStr := m.getSQLUpdate()
    condStr := m.getSQLCond()
    sql := "UPDATE "+ m.Table +" SET "+ updateStr + condStr +";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecInsert() (sql.Result, error) {
    insert := m.getSQLInsert()
    sql := "INSERT INTO "+m.Table+" "+insert+";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecReplace() (sql.Result, error) {
    insert := m.getSQLInsert()
    sql := "REPLACE INTO "+m.Table+" "+insert+";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecDelete() (sql.Result, error) {
    condStr := m.getSQLCond()
    sql := "DELETE FROM "+ m.Table + condStr +";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func GetLastInsertId(result sql.Result)(int64, error){
    return result.LastInsertId()
}

func GetRowsAffected(result sql.Result)(int64, error){
    return result.RowsAffected()
}

func (m *Model)getSQLCond() string {
    if str, ok := m.Cond.(string); ok{
        return str
    }
    var strArr []string
    if arr, ok := m.Cond.(map[string]string); ok{
        var strArr []string
        for k, v := range arr{
            strArr = append(strArr, k+"='"+v+"'")
        }
        return " WHERE "+ strings.Join(strArr, " && ")
    }
    if arr, ok := m.Cond.(map[string]interface{}); ok{
        for k, v := range arr{
            if isStr, ok := v.(string); ok{
                strArr = append(strArr, k+"='"+isStr+"'")
            }
            if isStrArr, ok := v.([]string); ok{
                for k, v := range isStrArr{
                    isStrArr[k] = "'"+v+"'"
                }
                strArr = append(strArr, k+" in ("+strings.Join(isStrArr, ",")+")")
            }
        }
        return " WHERE "+ strings.Join(strArr, " && ")
    }
    return ""
}

func (m *Model)getSQLField() string {
    if m.Field != ""{
        return m.Field
    }
    return "*"
}

func (m *Model)getSort() string {
    if m.Sort != ""{
        return " ORDER BY "+m.Sort+" "
    }
    return ""
}

func (m *Model)getSQLUpdate() string {
    var strArr []string
    for k, v := range m.Update{
        strArr = append(strArr, k+"='"+v+"'")
    }
    return strings.Join(strArr, ",")
}

func (m *Model)getSQLInsert() string {
    var fieldArr, valueArr []string
    for k, v := range m.Insert{
        fieldArr = append(fieldArr, k)
        valueArr = append(valueArr, "'"+v+"'")
    }
    return "("+ strings.Join(fieldArr, ",") +") values ("+ strings.Join(valueArr, ",") +")"
}

func (m *Model)getSQLFatch() int {
    if(m.Fetch == 0){
        return 0
    }
    return 1
}

func (m *Model)getSQLLimite() string {
    if strArr, ok := m.Limit.([2]int); ok{
        return " LIMIT "+fmt.Sprintf("%d", strArr[0])+", "+fmt.Sprintf("%d", strArr[1])
    }
    return ""
}

func (m *Model)Clean() {
    m.Cond = nil
    m.Insert = nil
    m.Update = nil
    m.Field = ""
    m.Table = ""
    m.Index = ""
    m.Limit = nil
    m.Sort = ""
    m.Fetch = 0
}<commit_msg>fix bugs<commit_after>package QesyDb

import(
    "database/sql"
    "fmt"
    "strings"
    "log"
    _"github.com/go-sql-driver/mysql"
    "reflect"
    "strconv"
    "time"
)

var Db *sql.DB

type Model struct{
    Cond interface{}
    Insert map[string]string
    Update map[string]string
    Field string
    Table string
    Index string
    Limit interface{}
    Sort string
    Fetch int
}

// Connect  is a method with a sql.
func Connect(connStr string){
    sqlDb, err := sql.Open("mysql", connStr)
    if(err != nil){
        log.Fatal("mysql connect error")
    }
    err = sqlDb.Ping()
    if(err != nil){
        log.Fatal("mysql ping error")
    }
	fmt.Println("mysql connect sueccss")
    Db = sqlDb
}

// ExecSelectIndex  is a method with a sql.
func (m *Model) ExecSelectIndex() (map[string]map[string]string, error) {
    resultsSlice, _ := m.ExecSelect()
    retArr := map[string]map[string]string{}
    for _, v := range resultsSlice{
        if v[m.Index] == "" {
            continue
        }
        retArr[v[m.Index]] = v
    }
    m.Clean()
    return retArr, nil
}

// ExecSelect is a method with a sql.
func (m *Model) ExecSelect() ([]map[string]string, error) {
    cond := m.getSQLCond()
    field := m.getSQLField()
    sort := m.getSort()
    limit := m.getSQLLimite()
    sql := "SELECT "+ field +" FROM "+m.Table+cond+sort+limit+";"
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    rows, err := stmt.Query()
    if(err != nil){
        return nil, err
    }
    defer rows.Close()
    fields, err := rows.Columns()
    if(err != nil){
        return nil, err
    }
    var resultsSlice []map[string]string
    for rows.Next(){
        result := make(map[string]string)
        var scanResultContainers []interface{}
        for i := 0; i < len(fields); i++ {
            var scanResultContainer interface{}
			scanResultContainers = append(scanResultContainers, &scanResultContainer)
        }
        if err := rows.Scan(scanResultContainers...); err != nil {
			return nil, err
		}
        for k, v := range fields {
            rawValue := reflect.Indirect(reflect.ValueOf(scanResultContainers[k]))
            if rawValue.Interface() == nil{
                continue
            }
            rawType := reflect.TypeOf(rawValue.Interface())
			rawVal := reflect.ValueOf(rawValue.Interface())
            var str string
            switch rawType.Kind() {
                case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
                    str = strconv.FormatInt(rawVal.Int(), 10)
                    result[v] = str
                case reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
                    str = strconv.FormatUint(rawVal.Uint(), 10)
                    result[v] = str
                case reflect.Float32, reflect.Float64:
                    str = strconv.FormatFloat(rawVal.Float(), 'f', -1, 64)
                    result[v] = str
                case reflect.Slice:
                    if rawType.Elem().Kind() == reflect.Uint8 {
                        result[v] = string(rawVal.Interface().([]byte))
                        break
                    }
                case reflect.String:
                    str = rawVal.String()
                    result[v] = str
                case reflect.Struct:
                    str = rawVal.Interface().(time.Time).Format("2006-01-02 15:04:05.000 -0700")
                    result[v] = str
                case reflect.Bool:
                    if rawVal.Bool() {
                        result[v] = "1"
                    } else {
                        result[v] = "0"
                    }
                }
            }
            resultsSlice = append(resultsSlice, result)
        }
        m.Clean()
        return resultsSlice, nil
}

func (m *Model) ExecSelectOne() (map[string]string, error) {
    resultsSlice, err := m.ExecSelect()
    if resultsSlice == nil{
        return resultsSlice, err
    }
    return resultsSlice[0], nil
}

func (m *Model) ExecUpdate() (sql.Result, error){
    updateStr := m.getSQLUpdate()
    condStr := m.getSQLCond()
    sql := "UPDATE "+ m.Table +" SET "+ updateStr + condStr +";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecInsert() (sql.Result, error) {
    insert := m.getSQLInsert()
    sql := "INSERT INTO "+m.Table+" "+insert+";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecReplace() (sql.Result, error) {
    insert := m.getSQLInsert()
    sql := "REPLACE INTO "+m.Table+" "+insert+";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func (m *Model) ExecDelete() (sql.Result, error) {
    condStr := m.getSQLCond()
    sql := "DELETE FROM "+ m.Table + condStr +";";
    stmt, err := Db.Prepare(sql)
    if(err != nil){
        return nil, err
    }
    result, err := stmt.Exec()
    m.Clean()
    return result, err
}

func GetLastInsertId(result sql.Result)(int64, error){
    return result.LastInsertId()
}

func GetRowsAffected(result sql.Result)(int64, error){
    return result.RowsAffected()
}

func (m *Model)getSQLCond() string {
    if str, ok := m.Cond.(string); ok{
        return str
    }
    var strArr []string
    if arr, ok := m.Cond.(map[string]string); ok{
        var strArr []string
        for k, v := range arr{
            strArr = append(strArr, k+"='"+v+"'")
        }
        return " WHERE "+ strings.Join(strArr, " && ")
    }
    if arr, ok := m.Cond.(map[string]interface{}); ok{
        for k, v := range arr{
            if isStr, ok := v.(string); ok{
                strArr = append(strArr, k+"='"+isStr+"'")
            }
            if isStrArr, ok := v.([]string); ok{
                for k, v := range isStrArr{
                    isStrArr[k] = "'"+v+"'"
                }
                strArr = append(strArr, k+" in ("+strings.Join(isStrArr, ",")+")")
            }
        }
        return " WHERE "+ strings.Join(strArr, " && ")
    }
    return ""
}

func (m *Model)getSQLField() string {
    if m.Field != ""{
        return m.Field
    }
    return "*"
}

func (m *Model)getSort() string {
    if m.Sort != ""{
        return " ORDER BY "+m.Sort+" "
    }
    return ""
}

func (m *Model)getSQLUpdate() string {
    var strArr []string
    for k, v := range m.Update{
        strArr = append(strArr, k+"='"+v+"'")
    }
    return strings.Join(strArr, ",")
}

func (m *Model)getSQLInsert() string {
    var fieldArr, valueArr []string
    for k, v := range m.Insert{
        fieldArr = append(fieldArr, k)
        valueArr = append(valueArr, "'"+v+"'")
    }
    return "("+ strings.Join(fieldArr, ",") +") values ("+ strings.Join(valueArr, ",") +")"
}

func (m *Model)getSQLFatch() int {
    if(m.Fetch == 0){
        return 0
    }
    return 1
}

func (m *Model)getSQLLimite() string {
    if strArr, ok := m.Limit.([2]int); ok{
        return " LIMIT "+fmt.Sprintf("%d", strArr[0])+", "+fmt.Sprintf("%d", strArr[1])
    }
    return ""
}

func (m *Model)Clean() {
    m.Cond = nil
    m.Insert = nil
    m.Update = nil
    m.Field = ""
    m.Table = ""
    m.Index = ""
    m.Limit = nil
    m.Sort = ""
    m.Fetch = 0
}<|endoftext|>


<commit_before>package main

import (
  "fmt"
)

func main() {
  fmt.Println("Usage: given an org, print list of users *without* 2fa enabled")
}
<commit_msg>Almost there: Only owners can use this filter<commit_after>package main

import (
  "fmt"
  "os"
  "github.com/octokit/go-octokit/octokit"
)

type Member struct {
  Login string `json:"login"`
}

func main() {
  if len(os.Args[1:]) < 1 {
    fmt.Println("Usage: given an org, print list of users *without* 2fa enabled")
    return
  }

  orgName := os.Args[1]
  fmt.Printf("Org Name is: %v", orgName)

  memberURL := octokit.Hyperlink("/orgs/{org}/members{?filter}") // type,page,per_page,sort
  url, err := memberURL.Expand(octokit.M{"org": orgName, "filter": "2fa_disabled"})

  fmt.Println("What is URL: " + url.String())

  if err != nil {
    fmt.Println("There was an error: ", err)
    return
  }
  fmt.Printf("URL is: %v\n", url)

  client := octokit.NewClient(nil)
  req, _ := client.NewRequest(url.String())
  var members []Member
  // result := client.get(url, &members)
  req.Get(&members)

  // if result.HasError() {
  //   fmt.Println("Error: ", err)
  //   return
  // }

  if len(members) == 0 { fmt.Println("No users have 2fa disabled.") }
  for i:=0; i<len(members); i++ {
    fmt.Println(members[i].Login)
  }
}
<|endoftext|>