<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 main

import (
	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/logger"
	"github.com/trivago/tgo"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"
)

const (
	coordinatorStateConfigure      = coordinatorState(iota)
	coordinatorStateStartProducers = coordinatorState(iota)
	coordinatorStateStartConsumers = coordinatorState(iota)
	coordinatorStateRunning        = coordinatorState(iota)
	coordinatorStateShutdown       = coordinatorState(iota)
	coordinatorStateStopConsumers  = coordinatorState(iota)
	coordinatorStateStopProducers  = coordinatorState(iota)
	coordinatorStateStopped        = coordinatorState(iota)
)

const (
	signalNone = signalType(iota)
	signalExit = signalType(iota)
	signalRoll = signalType(iota)
)

type coordinatorState byte
type signalType byte

// Coordinator is the main gollum instance taking care of starting and stopping
// plugins.
type Coordinator struct {
	consumers      []core.Consumer
	producers      []core.Producer
	routers        []core.Router
	consumerWorker *sync.WaitGroup
	producerWorker *sync.WaitGroup
	logConsumer    *core.LogConsumer
	state          coordinatorState
	signal         chan os.Signal
}

// NewCoordinator creates a new multplexer
func NewCoordinator() Coordinator {
	return Coordinator{
		consumerWorker: new(sync.WaitGroup),
		producerWorker: new(sync.WaitGroup),
		state:          coordinatorStateConfigure,
	}
}

// Configure processes the config and instantiates all valid plugins
func (co *Coordinator) Configure(conf *core.Config) error {
	// Make sure the log is printed to the fallback device if we are stuck here
	logFallback := time.AfterFunc(time.Duration(3)*time.Second, func() {
		logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
		logrusHookBuffer.Purge()
	})
	defer logFallback.Stop()

	// Initialize the plugins in the order of routers > producers > consumers
	// to match the order of reference between the different types.
	errors := tgo.NewErrorStack()
	errors.SetFormat(tgo.ErrorStackFormatCSV)

	if !co.configureRouters(conf) {
		errors.Pushf("At least one router failed to be configured")
	}
	if !co.configureProducers(conf) {
		errors.Pushf("At least one producer failed to be configured")
	}
	if !co.configureConsumers(conf) {
		errors.Pushf("At least one consumer failed to be configured")
	}
	if len(co.producers) == 0 {
		errors.Pushf("No valid producers found")
	}
	if len(co.consumers) <= 1 {
		errors.Pushf("No valid consumers found")
	}

	// As consumers might create new fallback router this is the first position
	// where we can add the wildcard producers to all streams. No new routers
	// created beyond this point must use StreamRegistry.AddWildcardProducersToRouter.

	core.StreamRegistry.AddAllWildcardProducersToAllRouters()
	return errors.OrNil()
}

// StartPlugins starts all plugins in the correct order.
func (co *Coordinator) StartPlugins() {
	// Launch routers
	for _, router := range co.routers {
		logrus.Debug("Starting ", reflect.TypeOf(router))
		if err := router.Start(); err != nil {
			logrus.WithError(err).Errorf("Failed to start router of type '%s'", reflect.TypeOf(router))
		}
	}

	// Launch producers
	co.state = coordinatorStateStartProducers
	for _, producer := range co.producers {
		producer := producer
		go tgo.WithRecoverShutdown(func() {
			logrus.Debug("Starting ", reflect.TypeOf(producer))
			producer.Produce(co.producerWorker)
		})
	}

	// Set final log target and purge the intermediate buffer
	if core.StreamRegistry.IsStreamRegistered(core.LogInternalStreamID) {
		// The _GOLLUM_ stream has listeners, so use LogConsumer to write to it
		if *flagLogColors != "always" {
			logrus.SetFormatter(logger.NewConsoleFormatter())
		}
		logrusHookBuffer.SetTargetHook(co.logConsumer)
		logrusHookBuffer.Purge()

	} else {
		logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
		logrusHookBuffer.Purge()
	}

	// Launch consumers
	co.state = coordinatorStateStartConsumers
	for _, consumer := range co.consumers {
		consumer := consumer
		go tgo.WithRecoverShutdown(func() {
			logrus.Debug("Starting ", reflect.TypeOf(consumer))
			consumer.Consume(co.consumerWorker)
		})
	}
}

// Run is essentially the Coordinator main loop.
// It listens for shutdown signals and updates global metrics
func (co *Coordinator) Run() {
	co.signal = newSignalHandler()
	defer signal.Stop(co.signal)

	logrus.Info("We be nice to them, if they be nice to us. (startup)")

	for {
		sig := <-co.signal
		switch translateSignal(sig) {
		case signalExit:
			logrus.Info("Master betrayed us. Wicked. Tricksy, False. (signal)")
			return // ### return, exit requested ###

		case signalRoll:
			for _, consumer := range co.consumers {
				consumer.Control() <- core.PluginControlRoll
			}
			for _, producer := range co.producers {
				producer.Control() <- core.PluginControlRoll
			}

		default:
		}
	}
}

// Shutdown all consumers and producers in a clean way.
// The internal log is flushed after the consumers have been shut down so that
// consumer related messages are still in the tlog.
// Producers are flushed after flushing the log, so producer related shutdown
// messages will be posted to stdout
func (co *Coordinator) Shutdown() {
	logrus.Info("Filthy little hobbites. They stole it from us. (shutdown)")

	stateAtShutdown := co.state
	co.state = coordinatorStateShutdown

	co.shutdownConsumers(stateAtShutdown)

	// Make sure remaining warning / errors are written to stderr
	logrus.Info("I'm not listening... I'm not listening... (flushing)")
	logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
	logrusHookBuffer.SetTargetHook(nil)
	logrusHookBuffer.Purge()

	// Shutdown producers
	co.shutdownProducers(stateAtShutdown)

	co.state = coordinatorStateStopped
}

func (co *Coordinator) configureRouters(conf *core.Config) bool {
	allFine := true
	routerConfigs := conf.GetRouters()
	for _, config := range routerConfigs {
		if _, hasStreams := config.Settings.Value("Stream"); !hasStreams {
			logrus.Errorf("Router '%s' has no stream set", config.ID)
			allFine = false
			continue // ### continue ###
		}

		logrus.Debugf("Instantiating router '%s'", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate router '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		routerPlugin := plugin.(core.Router)
		co.routers = append(co.routers, routerPlugin)

		logrus.Debugf("Instantiated '%s' (%s) as '%s'", config.ID, core.StreamRegistry.GetStreamName(routerPlugin.GetStreamID()), config.Typename)
		core.StreamRegistry.Register(routerPlugin, routerPlugin.GetStreamID())
	}

	return allFine
}

func (co *Coordinator) configureProducers(conf *core.Config) bool {
	co.state = coordinatorStateStartProducers
	allFine := true

	// All producers are added to the wildcard stream so that consumers can send
	// to all producers if required. The wildcard producer list is required
	// to add producers listening to all routers to all streams that are used.
	wildcardStream := core.StreamRegistry.GetRouterOrFallback(core.WildcardStreamID)
	producerConfigs := conf.GetProducers()

	for _, config := range producerConfigs {
		if _, hasStreams := config.Settings.Value("Streams"); !hasStreams {
			logrus.Errorf("Producer '%s' has no streams set", config.ID)
			allFine = false
			continue // ### continue ###
		}

		logrus.Debug("Instantiating ", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate producer '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		producer, _ := plugin.(core.Producer)
		co.producers = append(co.producers, producer)
		core.CountProducers()

		// Attach producer to streams
		streams := producer.Streams()
		for _, streamID := range streams {
			if streamID == core.WildcardStreamID {
				core.StreamRegistry.RegisterWildcardProducer(producer)
			} else {
				router := core.StreamRegistry.GetRouterOrFallback(streamID)
				router.AddProducer(producer)
			}
		}

		// Add producer to wildcard stream unless it only listens to internal streams
	searchinternal:
		for _, streamID := range streams {
			switch streamID {
			case core.LogInternalStreamID:
			default:
				wildcardStream.AddProducer(producer)
				break searchinternal
			}
		}
	}

	return allFine
}

func (co *Coordinator) configureConsumers(conf *core.Config) bool {
	co.state = coordinatorStateStartConsumers
	allFine := co.configureLogConsumer()

	consumerConfigs := conf.GetConsumers()
	for _, config := range consumerConfigs {
		if _, hasStreams := config.Settings.Value("Streams"); !hasStreams {
			logrus.Errorf("Consumer '%s' has no streams set", config.ID)
			allFine = false
			continue
		}

		logrus.Debug("Instantiating ", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate producer '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		consumer, _ := plugin.(core.Consumer)
		co.consumers = append(co.consumers, consumer)
		core.CountConsumers()
	}

	return allFine
}

func (co *Coordinator) configureLogConsumer() bool {
	config := core.NewPluginConfig("", "core.LogConsumer")
	configReader := core.NewPluginConfigReader(&config)

	co.logConsumer = new(core.LogConsumer)
	co.logConsumer.Configure(configReader)
	if configReader.Errors.Len() == 0 {
		co.consumers = append(co.consumers, co.logConsumer)
		return true
	}
	return false
}

func (co *Coordinator) shutdownConsumers(stateAtShutdown coordinatorState) {
	if stateAtShutdown >= coordinatorStateStartConsumers {
		co.state = coordinatorStateStopConsumers
		waitTimeout := time.Duration(0)

		logrus.Debug("Telling consumers to stop")
		for _, cons := range co.consumers {
			timeout := cons.GetShutdownTimeout()
			if timeout > waitTimeout {
				waitTimeout = timeout
			}

			// Skip log consumer so we get clean log messages for all consumers
			if cons != co.logConsumer {
				cons.Control() <- core.PluginControlStopConsumer
			}
		}

		waitTimeout *= 10
		logrus.Debugf("Waiting for consumers to stop. Forced shutdown after %.2f seconds.", waitTimeout.Seconds())
		if co.logConsumer != nil {
			co.logConsumer.Control() <- core.PluginControlStopConsumer
			// No logs in _GOLLUM_ after this point
		}

		if !tgo.ReturnAfter(waitTimeout, co.consumerWorker.Wait) {
			logrus.Error("At least one consumer found to be blocking.")
		}
	}
}

func (co *Coordinator) shutdownProducers(stateAtShutdown coordinatorState) {
	if stateAtShutdown >= coordinatorStateStartProducers {
		co.state = coordinatorStateStopProducers
		waitTimeout := time.Duration(0)

		logrus.Debug("Telling producers to stop")
		for _, prod := range co.producers {
			timeout := prod.GetShutdownTimeout()
			if timeout > waitTimeout {
				waitTimeout = timeout
			}
			prod.Control() <- core.PluginControlStopProducer
		}

		waitTimeout *= 10
		logrus.Debugf("Waiting for producers to stop. Forced shutdown after %.2f seconds.", waitTimeout.Seconds())
		if !tgo.ReturnAfter(waitTimeout, co.producerWorker.Wait) {
			logrus.Error("At least one producer found to be blocking.")
		}
	}
}
<commit_msg>fix colors flag<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 main

import (
	"os"
	"os/signal"
	"reflect"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/logger"
	"github.com/trivago/tgo"
)

const (
	coordinatorStateConfigure      = coordinatorState(iota)
	coordinatorStateStartProducers = coordinatorState(iota)
	coordinatorStateStartConsumers = coordinatorState(iota)
	coordinatorStateRunning        = coordinatorState(iota)
	coordinatorStateShutdown       = coordinatorState(iota)
	coordinatorStateStopConsumers  = coordinatorState(iota)
	coordinatorStateStopProducers  = coordinatorState(iota)
	coordinatorStateStopped        = coordinatorState(iota)
)

const (
	signalNone = signalType(iota)
	signalExit = signalType(iota)
	signalRoll = signalType(iota)
)

type coordinatorState byte
type signalType byte

// Coordinator is the main gollum instance taking care of starting and stopping
// plugins.
type Coordinator struct {
	consumers      []core.Consumer
	producers      []core.Producer
	routers        []core.Router
	consumerWorker *sync.WaitGroup
	producerWorker *sync.WaitGroup
	logConsumer    *core.LogConsumer
	state          coordinatorState
	signal         chan os.Signal
}

// NewCoordinator creates a new multplexer
func NewCoordinator() Coordinator {
	return Coordinator{
		consumerWorker: new(sync.WaitGroup),
		producerWorker: new(sync.WaitGroup),
		state:          coordinatorStateConfigure,
	}
}

// Configure processes the config and instantiates all valid plugins
func (co *Coordinator) Configure(conf *core.Config) error {
	// Make sure the log is printed to the fallback device if we are stuck here
	logFallback := time.AfterFunc(time.Duration(3)*time.Second, func() {
		logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
		logrusHookBuffer.Purge()
	})
	defer logFallback.Stop()

	// Initialize the plugins in the order of routers > producers > consumers
	// to match the order of reference between the different types.
	errors := tgo.NewErrorStack()
	errors.SetFormat(tgo.ErrorStackFormatCSV)

	if !co.configureRouters(conf) {
		errors.Pushf("At least one router failed to be configured")
	}
	if !co.configureProducers(conf) {
		errors.Pushf("At least one producer failed to be configured")
	}
	if !co.configureConsumers(conf) {
		errors.Pushf("At least one consumer failed to be configured")
	}
	if len(co.producers) == 0 {
		errors.Pushf("No valid producers found")
	}
	if len(co.consumers) <= 1 {
		errors.Pushf("No valid consumers found")
	}

	// As consumers might create new fallback router this is the first position
	// where we can add the wildcard producers to all streams. No new routers
	// created beyond this point must use StreamRegistry.AddWildcardProducersToRouter.

	core.StreamRegistry.AddAllWildcardProducersToAllRouters()
	return errors.OrNil()
}

// StartPlugins starts all plugins in the correct order.
func (co *Coordinator) StartPlugins() {
	// Launch routers
	for _, router := range co.routers {
		logrus.Debug("Starting ", reflect.TypeOf(router))
		if err := router.Start(); err != nil {
			logrus.WithError(err).Errorf("Failed to start router of type '%s'", reflect.TypeOf(router))
		}
	}

	// Launch producers
	co.state = coordinatorStateStartProducers
	for _, producer := range co.producers {
		producer := producer
		go tgo.WithRecoverShutdown(func() {
			logrus.Debug("Starting ", reflect.TypeOf(producer))
			producer.Produce(co.producerWorker)
		})
	}

	// Set final log target and purge the intermediate buffer
	if core.StreamRegistry.IsStreamRegistered(core.LogInternalStreamID) {
		// The _GOLLUM_ stream has listeners, so use LogConsumer to write to it
		if *flagLogColors == "always" {
			logrus.SetFormatter(logger.NewConsoleFormatter())
		}
		logrusHookBuffer.SetTargetHook(co.logConsumer)
		logrusHookBuffer.Purge()

	} else {
		logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
		logrusHookBuffer.Purge()
	}

	// Launch consumers
	co.state = coordinatorStateStartConsumers
	for _, consumer := range co.consumers {
		consumer := consumer
		go tgo.WithRecoverShutdown(func() {
			logrus.Debug("Starting ", reflect.TypeOf(consumer))
			consumer.Consume(co.consumerWorker)
		})
	}
}

// Run is essentially the Coordinator main loop.
// It listens for shutdown signals and updates global metrics
func (co *Coordinator) Run() {
	co.signal = newSignalHandler()
	defer signal.Stop(co.signal)

	logrus.Info("We be nice to them, if they be nice to us. (startup)")

	for {
		sig := <-co.signal
		switch translateSignal(sig) {
		case signalExit:
			logrus.Info("Master betrayed us. Wicked. Tricksy, False. (signal)")
			return // ### return, exit requested ###

		case signalRoll:
			for _, consumer := range co.consumers {
				consumer.Control() <- core.PluginControlRoll
			}
			for _, producer := range co.producers {
				producer.Control() <- core.PluginControlRoll
			}

		default:
		}
	}
}

// Shutdown all consumers and producers in a clean way.
// The internal log is flushed after the consumers have been shut down so that
// consumer related messages are still in the tlog.
// Producers are flushed after flushing the log, so producer related shutdown
// messages will be posted to stdout
func (co *Coordinator) Shutdown() {
	logrus.Info("Filthy little hobbites. They stole it from us. (shutdown)")

	stateAtShutdown := co.state
	co.state = coordinatorStateShutdown

	co.shutdownConsumers(stateAtShutdown)

	// Make sure remaining warning / errors are written to stderr
	logrus.Info("I'm not listening... I'm not listening... (flushing)")
	logrusHookBuffer.SetTargetWriter(logger.FallbackLogDevice)
	logrusHookBuffer.SetTargetHook(nil)
	logrusHookBuffer.Purge()

	// Shutdown producers
	co.shutdownProducers(stateAtShutdown)

	co.state = coordinatorStateStopped
}

func (co *Coordinator) configureRouters(conf *core.Config) bool {
	allFine := true
	routerConfigs := conf.GetRouters()
	for _, config := range routerConfigs {
		if _, hasStreams := config.Settings.Value("Stream"); !hasStreams {
			logrus.Errorf("Router '%s' has no stream set", config.ID)
			allFine = false
			continue // ### continue ###
		}

		logrus.Debugf("Instantiating router '%s'", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate router '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		routerPlugin := plugin.(core.Router)
		co.routers = append(co.routers, routerPlugin)

		logrus.Debugf("Instantiated '%s' (%s) as '%s'", config.ID, core.StreamRegistry.GetStreamName(routerPlugin.GetStreamID()), config.Typename)
		core.StreamRegistry.Register(routerPlugin, routerPlugin.GetStreamID())
	}

	return allFine
}

func (co *Coordinator) configureProducers(conf *core.Config) bool {
	co.state = coordinatorStateStartProducers
	allFine := true

	// All producers are added to the wildcard stream so that consumers can send
	// to all producers if required. The wildcard producer list is required
	// to add producers listening to all routers to all streams that are used.
	wildcardStream := core.StreamRegistry.GetRouterOrFallback(core.WildcardStreamID)
	producerConfigs := conf.GetProducers()

	for _, config := range producerConfigs {
		if _, hasStreams := config.Settings.Value("Streams"); !hasStreams {
			logrus.Errorf("Producer '%s' has no streams set", config.ID)
			allFine = false
			continue // ### continue ###
		}

		logrus.Debug("Instantiating ", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate producer '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		producer, _ := plugin.(core.Producer)
		co.producers = append(co.producers, producer)
		core.CountProducers()

		// Attach producer to streams
		streams := producer.Streams()
		for _, streamID := range streams {
			if streamID == core.WildcardStreamID {
				core.StreamRegistry.RegisterWildcardProducer(producer)
			} else {
				router := core.StreamRegistry.GetRouterOrFallback(streamID)
				router.AddProducer(producer)
			}
		}

		// Add producer to wildcard stream unless it only listens to internal streams
	searchinternal:
		for _, streamID := range streams {
			switch streamID {
			case core.LogInternalStreamID:
			default:
				wildcardStream.AddProducer(producer)
				break searchinternal
			}
		}
	}

	return allFine
}

func (co *Coordinator) configureConsumers(conf *core.Config) bool {
	co.state = coordinatorStateStartConsumers
	allFine := co.configureLogConsumer()

	consumerConfigs := conf.GetConsumers()
	for _, config := range consumerConfigs {
		if _, hasStreams := config.Settings.Value("Streams"); !hasStreams {
			logrus.Errorf("Consumer '%s' has no streams set", config.ID)
			allFine = false
			continue
		}

		logrus.Debug("Instantiating ", config.ID)
		plugin, err := core.NewPluginWithConfig(config)
		if err != nil {
			logrus.WithError(err).Errorf("Failed to instantiate producer '%s'", config.ID)
			allFine = false
			continue // ### continue ###
		}

		consumer, _ := plugin.(core.Consumer)
		co.consumers = append(co.consumers, consumer)
		core.CountConsumers()
	}

	return allFine
}

func (co *Coordinator) configureLogConsumer() bool {
	config := core.NewPluginConfig("", "core.LogConsumer")
	configReader := core.NewPluginConfigReader(&config)

	co.logConsumer = new(core.LogConsumer)
	co.logConsumer.Configure(configReader)
	if configReader.Errors.Len() == 0 {
		co.consumers = append(co.consumers, co.logConsumer)
		return true
	}
	return false
}

func (co *Coordinator) shutdownConsumers(stateAtShutdown coordinatorState) {
	if stateAtShutdown >= coordinatorStateStartConsumers {
		co.state = coordinatorStateStopConsumers
		waitTimeout := time.Duration(0)

		logrus.Debug("Telling consumers to stop")
		for _, cons := range co.consumers {
			timeout := cons.GetShutdownTimeout()
			if timeout > waitTimeout {
				waitTimeout = timeout
			}

			// Skip log consumer so we get clean log messages for all consumers
			if cons != co.logConsumer {
				cons.Control() <- core.PluginControlStopConsumer
			}
		}

		waitTimeout *= 10
		logrus.Debugf("Waiting for consumers to stop. Forced shutdown after %.2f seconds.", waitTimeout.Seconds())
		if co.logConsumer != nil {
			co.logConsumer.Control() <- core.PluginControlStopConsumer
			// No logs in _GOLLUM_ after this point
		}

		if !tgo.ReturnAfter(waitTimeout, co.consumerWorker.Wait) {
			logrus.Error("At least one consumer found to be blocking.")
		}
	}
}

func (co *Coordinator) shutdownProducers(stateAtShutdown coordinatorState) {
	if stateAtShutdown >= coordinatorStateStartProducers {
		co.state = coordinatorStateStopProducers
		waitTimeout := time.Duration(0)

		logrus.Debug("Telling producers to stop")
		for _, prod := range co.producers {
			timeout := prod.GetShutdownTimeout()
			if timeout > waitTimeout {
				waitTimeout = timeout
			}
			prod.Control() <- core.PluginControlStopProducer
		}

		waitTimeout *= 10
		logrus.Debugf("Waiting for producers to stop. Forced shutdown after %.2f seconds.", waitTimeout.Seconds())
		if !tgo.ReturnAfter(waitTimeout, co.producerWorker.Wait) {
			logrus.Error("At least one producer found to be blocking.")
		}
	}
}
<|endoftext|>


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

package maas

import (
	"bytes"
	"io/ioutil"

	"github.com/juju/errors"
	"github.com/juju/gomaasapi"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type maas2StorageSuite struct {
	maas2Suite
}

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

func (s *maas2StorageSuite) makeStorage(c *gc.C, controller gomaasapi.Controller) *maas2Storage {
	storage, ok := NewStorage(s.makeEnviron(c, controller)).(*maas2Storage)
	c.Assert(ok, jc.IsTrue)
	ecfg := storage.environ.ecfg()
	ecfg.attrs["maas-agent-name"] = "prefix"
	return storage
}

func getController(c *gc.C, storage *maas2Storage) *fakeController {
	result, ok := storage.maasController.(*fakeController)
	c.Assert(ok, jc.IsTrue)
	return result
}

func (s *maas2StorageSuite) TestGetNoSuchFile(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("This file no existence"),
	})
	_, err := storage.Get("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "This file no existence")
}

func (s *maas2StorageSuite) TestGetReadFails(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		files: []gomaasapi.File{
			&fakeFile{
				name:  "prefix-grasshopper.avi",
				error: errors.New("read error"),
			},
		},
	})
	_, err := storage.Get("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "read error")
}

func (s *maas2StorageSuite) TestGetSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{
				name:     "prefix-grasshopper.avi",
				contents: []byte("The man in the high castle"),
			},
		},
	}
	storage := s.makeStorage(c, controller)
	reader, err := storage.Get("grasshopper.avi")
	c.Assert(err, jc.ErrorIsNil)
	defer reader.Close()
	result, err := ioutil.ReadAll(reader)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.DeepEquals, []byte("The man in the high castle"))
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-grasshopper.avi")
}

func (s *maas2StorageSuite) TestListError(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("couldn't list files"),
	})
	_, err := storage.List("american-territories")
	c.Assert(err, gc.ErrorMatches, "couldn't list files")
}

func (s *maas2StorageSuite) TestListSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-julianna"},
			&fakeFile{name: "prefix-frank"},
		},
	}
	storage := s.makeStorage(c, controller)
	result, err := storage.List("grasshopper")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, []string{"frank", "julianna"})
	c.Assert(controller.filesPrefix, gc.Equals, "prefix-grasshopper")
}

func (s *maas2StorageSuite) TestURLError(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("no such file"),
	})
	_, err := storage.URL("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "no such file")
}

func (s *maas2StorageSuite) TestURLSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-grasshopper.avi", url: "heavy lies"},
		},
	}
	storage := s.makeStorage(c, controller)
	result, err := storage.URL("grasshopper.avi")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.Equals, "heavy lies")
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-grasshopper.avi")
}

func (s *maas2StorageSuite) TestPut(c *gc.C) {
	controller := &fakeController{
		filesError: errors.New("oh no!"),
	}
	storage := s.makeStorage(c, controller)
	reader := bytes.NewReader([]byte{})
	err := storage.Put("riff", reader, 10)
	c.Assert(err, gc.ErrorMatches, "oh no!")
	c.Assert(controller.addFileArgs, jc.DeepEquals, gomaasapi.AddFileArgs{
		Filename: "prefix-riff",
		Reader:   reader,
		Length:   10,
	})
}

func (s *maas2StorageSuite) TestRemoveNoSuchFile(c *gc.C) {
	controller := &fakeController{
		filesError: errors.New("oh no!"),
	}
	storage := s.makeStorage(c, controller)
	err := storage.Remove("FIOS")
	c.Assert(err, gc.ErrorMatches, "oh no!")
}

func (s *maas2StorageSuite) TestRemoveErrorFromDelete(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-FIOS", error: errors.New("protected")},
		},
	}
	storage := s.makeStorage(c, controller)
	err := storage.Remove("FIOS")
	c.Assert(err, gc.ErrorMatches, "protected")
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-FIOS")
}

func (s *maas2StorageSuite) TestRemoveAll(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-zack"},
			&fakeFile{name: "prefix-kevin", error: errors.New("oops")},
			&fakeFile{name: "prefix-jim"},
			&fakeFile{name: "prefix-riff"},
		},
	}
	storage := s.makeStorage(c, controller)
	err := storage.RemoveAll()
	c.Assert(err, gc.ErrorMatches, "cannot delete all provider state: oops")
	c.Assert(controller.filesPrefix, gc.Equals, "prefix-")

	deleteds := make([]bool, 4)
	for i, file := range controller.files {
		deleteds[i] = file.(*fakeFile).deleted
	}
	c.Assert(deleteds, jc.DeepEquals, []bool{true, true, true, true})
}
<commit_msg>Test Get not found error<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package maas

import (
	"bytes"
	"io/ioutil"

	"github.com/juju/errors"
	"github.com/juju/gomaasapi"
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"
)

type maas2StorageSuite struct {
	maas2Suite
}

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

func (s *maas2StorageSuite) makeStorage(c *gc.C, controller gomaasapi.Controller) *maas2Storage {
	storage, ok := NewStorage(s.makeEnviron(c, controller)).(*maas2Storage)
	c.Assert(ok, jc.IsTrue)
	ecfg := storage.environ.ecfg()
	ecfg.attrs["maas-agent-name"] = "prefix"
	return storage
}

func getController(c *gc.C, storage *maas2Storage) *fakeController {
	result, ok := storage.maasController.(*fakeController)
	c.Assert(ok, jc.IsTrue)
	return result
}

func (s *maas2StorageSuite) TestGetNoSuchFile(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("This file no existence"),
	})
	_, err := storage.Get("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "This file no existence")
}

func (s *maas2StorageSuite) TestGetReadFails(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		files: []gomaasapi.File{
			&fakeFile{
				name:  "prefix-grasshopper.avi",
				error: errors.New("read error"),
			},
		},
	})
	_, err := storage.Get("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "read error")
}

func (s *maas2StorageSuite) TestGetNotFounc(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: gomaasapi.NewNoMatchError("wee"),
	})
	_, err := storage.Get("grasshopper.avi")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func (s *maas2StorageSuite) TestGetSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{
				name:     "prefix-grasshopper.avi",
				contents: []byte("The man in the high castle"),
			},
		},
	}
	storage := s.makeStorage(c, controller)
	reader, err := storage.Get("grasshopper.avi")
	c.Assert(err, jc.ErrorIsNil)
	defer reader.Close()
	result, err := ioutil.ReadAll(reader)
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.DeepEquals, []byte("The man in the high castle"))
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-grasshopper.avi")
}

func (s *maas2StorageSuite) TestListError(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("couldn't list files"),
	})
	_, err := storage.List("american-territories")
	c.Assert(err, gc.ErrorMatches, "couldn't list files")
}

func (s *maas2StorageSuite) TestListSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-julianna"},
			&fakeFile{name: "prefix-frank"},
		},
	}
	storage := s.makeStorage(c, controller)
	result, err := storage.List("grasshopper")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, jc.DeepEquals, []string{"frank", "julianna"})
	c.Assert(controller.filesPrefix, gc.Equals, "prefix-grasshopper")
}

func (s *maas2StorageSuite) TestURLError(c *gc.C) {
	storage := s.makeStorage(c, &fakeController{
		filesError: errors.New("no such file"),
	})
	_, err := storage.URL("grasshopper.avi")
	c.Assert(err, gc.ErrorMatches, "no such file")
}

func (s *maas2StorageSuite) TestURLSuccess(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-grasshopper.avi", url: "heavy lies"},
		},
	}
	storage := s.makeStorage(c, controller)
	result, err := storage.URL("grasshopper.avi")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(result, gc.Equals, "heavy lies")
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-grasshopper.avi")
}

func (s *maas2StorageSuite) TestPut(c *gc.C) {
	controller := &fakeController{
		filesError: errors.New("oh no!"),
	}
	storage := s.makeStorage(c, controller)
	reader := bytes.NewReader([]byte{})
	err := storage.Put("riff", reader, 10)
	c.Assert(err, gc.ErrorMatches, "oh no!")
	c.Assert(controller.addFileArgs, jc.DeepEquals, gomaasapi.AddFileArgs{
		Filename: "prefix-riff",
		Reader:   reader,
		Length:   10,
	})
}

func (s *maas2StorageSuite) TestRemoveNoSuchFile(c *gc.C) {
	controller := &fakeController{
		filesError: errors.New("oh no!"),
	}
	storage := s.makeStorage(c, controller)
	err := storage.Remove("FIOS")
	c.Assert(err, gc.ErrorMatches, "oh no!")
}

func (s *maas2StorageSuite) TestRemoveErrorFromDelete(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-FIOS", error: errors.New("protected")},
		},
	}
	storage := s.makeStorage(c, controller)
	err := storage.Remove("FIOS")
	c.Assert(err, gc.ErrorMatches, "protected")
	c.Assert(controller.getFileFilename, gc.Equals, "prefix-FIOS")
}

func (s *maas2StorageSuite) TestRemoveAll(c *gc.C) {
	controller := &fakeController{
		files: []gomaasapi.File{
			&fakeFile{name: "prefix-zack"},
			&fakeFile{name: "prefix-kevin", error: errors.New("oops")},
			&fakeFile{name: "prefix-jim"},
			&fakeFile{name: "prefix-riff"},
		},
	}
	storage := s.makeStorage(c, controller)
	err := storage.RemoveAll()
	c.Assert(err, gc.ErrorMatches, "cannot delete all provider state: oops")
	c.Assert(controller.filesPrefix, gc.Equals, "prefix-")

	deleteds := make([]bool, 4)
	for i, file := range controller.files {
		deleteds[i] = file.(*fakeFile).deleted
	}
	c.Assert(deleteds, jc.DeepEquals, []bool{true, true, true, true})
}
<|endoftext|>


<commit_before>package hclsyntax

import (
	"testing"

	"github.com/hashicorp/hcl/v2"
	"github.com/zclconf/go-cty/cty"
)

func TestTemplateExprParseAndValue(t *testing.T) {
	// This is a combo test that exercises both the parser and the Value
	// method, with the focus on the latter but indirectly testing the former.
	tests := []struct {
		input     string
		ctx       *hcl.EvalContext
		want      cty.Value
		diagCount int
	}{
		{
			`1`,
			nil,
			cty.StringVal("1"),
			0,
		},
		{
			`(1)`,
			nil,
			cty.StringVal("(1)"),
			0,
		},
		{
			`true`,
			nil,
			cty.StringVal("true"),
			0,
		},
		{
			`
hello world
`,
			nil,
			cty.StringVal("\nhello world\n"),
			0,
		},
		{
			`hello ${"world"}`,
			nil,
			cty.StringVal("hello world"),
			0,
		},
		{
			`hello\nworld`, // backslash escapes not supported in bare templates
			nil,
			cty.StringVal("hello\\nworld"),
			0,
		},
		{
			`hello ${12.5}`,
			nil,
			cty.StringVal("hello 12.5"),
			0,
		},
		{
			`silly ${"${"nesting"}"}`,
			nil,
			cty.StringVal("silly nesting"),
			0,
		},
		{
			`silly ${"${true}"}`,
			nil,
			cty.StringVal("silly true"),
			0,
		},
		{
			`hello $${escaped}`,
			nil,
			cty.StringVal("hello ${escaped}"),
			0,
		},
		{
			`hello $$nonescape`,
			nil,
			cty.StringVal("hello $$nonescape"),
			0,
		},
		{
			`hello %${"world"}`,
			nil,
			cty.StringVal("hello %world"),
			0,
		},
		{
			`${true}`,
			nil,
			cty.True, // any single expression is unwrapped without stringification
			0,
		},
		{
			`trim ${~ "trim"}`,
			nil,
			cty.StringVal("trimtrim"),
			0,
		},
		{
			`${"trim" ~} trim`,
			nil,
			cty.StringVal("trimtrim"),
			0,
		},
		{
			`trim
${~"trim"~}
trim`,
			nil,
			cty.StringVal("trimtrimtrim"),
			0,
		},
		{
			` ${~ true ~} `,
			nil,
			cty.StringVal("true"), // can't trim space to reduce to a single expression
			0,
		},
		{
			`${"hello "}${~"trim"~}${" hello"}`,
			nil,
			cty.StringVal("hello trim hello"), // trimming can't reach into a neighboring interpolation
			0,
		},
		{
			`${true}${~"trim"~}${true}`,
			nil,
			cty.StringVal("truetrimtrue"), // trimming is no-op of neighbors aren't literal strings
			0,
		},

		{
			`%{ if true ~} hello %{~ endif }`,
			nil,
			cty.StringVal("hello"),
			0,
		},
		{
			`%{ if false ~} hello %{~ endif}`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ if true ~} hello %{~ else ~} goodbye %{~ endif }`,
			nil,
			cty.StringVal("hello"),
			0,
		},
		{
			`%{ if false ~} hello %{~ else ~} goodbye %{~ endif }`,
			nil,
			cty.StringVal("goodbye"),
			0,
		},
		{
			`%{ if true ~} %{~ if false ~} hello %{~ else ~} goodbye %{~ endif ~} %{~ endif }`,
			nil,
			cty.StringVal("goodbye"),
			0,
		},
		{
			`%{ if false ~} %{~ if false ~} hello %{~ else ~} goodbye %{~ endif ~} %{~ endif }`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ of true ~} hello %{~ endif}`,
			nil,
			cty.UnknownVal(cty.String),
			2, // "of" is not a valid control keyword, and "endif" is therefore also unexpected
		},
		{
			`%{ for v in ["a", "b", "c"] }${v}%{ endfor }`,
			nil,
			cty.StringVal("abc"),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] } ${v} %{ endfor }`,
			nil,
			cty.StringVal(" a  b  c "),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] ~} ${v} %{~ endfor }`,
			nil,
			cty.StringVal("abc"),
			0,
		},
		{
			`%{ for v in [] }${v}%{ endfor }`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ for i, v in ["a", "b", "c"] }${i}${v}%{ endfor }`,
			nil,
			cty.StringVal("0a1b2c"),
			0,
		},
		{
			`%{ for k, v in {"A" = "a", "B" = "b", "C" = "c"} }${k}${v}%{ endfor }`,
			nil,
			cty.StringVal("AaBbCc"),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] }${v}${nl}%{ endfor }`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"nl": cty.StringVal("\n"),
				},
			},
			cty.StringVal("a\nb\nc\n"),
			0,
		},
		{
			`\n`, // backslash escapes are not interpreted in template literals
			nil,
			cty.StringVal("\\n"),
			0,
		},
		{
			`\uu1234`, // backslash escapes are not interpreted in template literals
			nil,       // (this is intentionally an invalid one to ensure we don't produce an error)
			cty.StringVal("\\uu1234"),
			0,
		},
		{
			`$`,
			nil,
			cty.StringVal("$"),
			0,
		},
		{
			`$$`,
			nil,
			cty.StringVal("$$"),
			0,
		},
		{
			`%`,
			nil,
			cty.StringVal("%"),
			0,
		},
		{
			`%%`,
			nil,
			cty.StringVal("%%"),
			0,
		},
		{
			`hello %%{ if true }world%%{ endif }`,
			nil,
			cty.StringVal(`hello %{ if true }world%{ endif }`),
			0,
		},
		{
			`hello $%{ if true }world%{ endif }`,
			nil,
			cty.StringVal("hello $world"),
			0,
		},
		{
			`%{ endif }`,
			nil,
			cty.UnknownVal(cty.String),
			1, // Unexpected endif directive
		},
		{
			`%{ endfor }`,
			nil,
			cty.UnknownVal(cty.String),
			1, // Unexpected endfor directive
		},
		{ // marks from uninterpolated values are ignored
			`hello%{ if false } ${target}%{ endif }`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"target": cty.StringVal("world").WithMarks(cty.NewValueMarks("sensitive")),
				},
			},
			cty.StringVal("hello"),
			0,
		},
		{ // marks from interpolated values are passed through
			`${greeting} ${target}`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"greeting": cty.StringVal("hello").WithMarks(cty.NewValueMarks("english")),
					"target":   cty.StringVal("world").WithMarks(cty.NewValueMarks("sensitive")),
				},
			},
			cty.StringVal("hello world").WithMarks(cty.NewValueMarks("english", "sensitive")),
			0,
		},
	}

	for _, test := range tests {
		t.Run(test.input, func(t *testing.T) {
			expr, parseDiags := ParseTemplate([]byte(test.input), "", hcl.Pos{Line: 1, Column: 1, Byte: 0})

			got, valDiags := expr.Value(test.ctx)

			diagCount := len(parseDiags) + len(valDiags)

			if diagCount != test.diagCount {
				t.Errorf("wrong number of diagnostics %d; want %d", diagCount, test.diagCount)
				for _, diag := range parseDiags {
					t.Logf(" - %s", diag.Error())
				}
				for _, diag := range valDiags {
					t.Logf(" - %s", diag.Error())
				}
			}

			if !got.RawEquals(test.want) {
				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.want)
			}
		})
	}

}
<commit_msg>hclsyntax: Test for interpolating from marked map<commit_after>package hclsyntax

import (
	"testing"

	"github.com/hashicorp/hcl/v2"
	"github.com/zclconf/go-cty/cty"
)

func TestTemplateExprParseAndValue(t *testing.T) {
	// This is a combo test that exercises both the parser and the Value
	// method, with the focus on the latter but indirectly testing the former.
	tests := []struct {
		input     string
		ctx       *hcl.EvalContext
		want      cty.Value
		diagCount int
	}{
		{
			`1`,
			nil,
			cty.StringVal("1"),
			0,
		},
		{
			`(1)`,
			nil,
			cty.StringVal("(1)"),
			0,
		},
		{
			`true`,
			nil,
			cty.StringVal("true"),
			0,
		},
		{
			`
hello world
`,
			nil,
			cty.StringVal("\nhello world\n"),
			0,
		},
		{
			`hello ${"world"}`,
			nil,
			cty.StringVal("hello world"),
			0,
		},
		{
			`hello\nworld`, // backslash escapes not supported in bare templates
			nil,
			cty.StringVal("hello\\nworld"),
			0,
		},
		{
			`hello ${12.5}`,
			nil,
			cty.StringVal("hello 12.5"),
			0,
		},
		{
			`silly ${"${"nesting"}"}`,
			nil,
			cty.StringVal("silly nesting"),
			0,
		},
		{
			`silly ${"${true}"}`,
			nil,
			cty.StringVal("silly true"),
			0,
		},
		{
			`hello $${escaped}`,
			nil,
			cty.StringVal("hello ${escaped}"),
			0,
		},
		{
			`hello $$nonescape`,
			nil,
			cty.StringVal("hello $$nonescape"),
			0,
		},
		{
			`hello %${"world"}`,
			nil,
			cty.StringVal("hello %world"),
			0,
		},
		{
			`${true}`,
			nil,
			cty.True, // any single expression is unwrapped without stringification
			0,
		},
		{
			`trim ${~ "trim"}`,
			nil,
			cty.StringVal("trimtrim"),
			0,
		},
		{
			`${"trim" ~} trim`,
			nil,
			cty.StringVal("trimtrim"),
			0,
		},
		{
			`trim
${~"trim"~}
trim`,
			nil,
			cty.StringVal("trimtrimtrim"),
			0,
		},
		{
			` ${~ true ~} `,
			nil,
			cty.StringVal("true"), // can't trim space to reduce to a single expression
			0,
		},
		{
			`${"hello "}${~"trim"~}${" hello"}`,
			nil,
			cty.StringVal("hello trim hello"), // trimming can't reach into a neighboring interpolation
			0,
		},
		{
			`${true}${~"trim"~}${true}`,
			nil,
			cty.StringVal("truetrimtrue"), // trimming is no-op of neighbors aren't literal strings
			0,
		},

		{
			`%{ if true ~} hello %{~ endif }`,
			nil,
			cty.StringVal("hello"),
			0,
		},
		{
			`%{ if false ~} hello %{~ endif}`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ if true ~} hello %{~ else ~} goodbye %{~ endif }`,
			nil,
			cty.StringVal("hello"),
			0,
		},
		{
			`%{ if false ~} hello %{~ else ~} goodbye %{~ endif }`,
			nil,
			cty.StringVal("goodbye"),
			0,
		},
		{
			`%{ if true ~} %{~ if false ~} hello %{~ else ~} goodbye %{~ endif ~} %{~ endif }`,
			nil,
			cty.StringVal("goodbye"),
			0,
		},
		{
			`%{ if false ~} %{~ if false ~} hello %{~ else ~} goodbye %{~ endif ~} %{~ endif }`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ of true ~} hello %{~ endif}`,
			nil,
			cty.UnknownVal(cty.String),
			2, // "of" is not a valid control keyword, and "endif" is therefore also unexpected
		},
		{
			`%{ for v in ["a", "b", "c"] }${v}%{ endfor }`,
			nil,
			cty.StringVal("abc"),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] } ${v} %{ endfor }`,
			nil,
			cty.StringVal(" a  b  c "),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] ~} ${v} %{~ endfor }`,
			nil,
			cty.StringVal("abc"),
			0,
		},
		{
			`%{ for v in [] }${v}%{ endfor }`,
			nil,
			cty.StringVal(""),
			0,
		},
		{
			`%{ for i, v in ["a", "b", "c"] }${i}${v}%{ endfor }`,
			nil,
			cty.StringVal("0a1b2c"),
			0,
		},
		{
			`%{ for k, v in {"A" = "a", "B" = "b", "C" = "c"} }${k}${v}%{ endfor }`,
			nil,
			cty.StringVal("AaBbCc"),
			0,
		},
		{
			`%{ for v in ["a", "b", "c"] }${v}${nl}%{ endfor }`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"nl": cty.StringVal("\n"),
				},
			},
			cty.StringVal("a\nb\nc\n"),
			0,
		},
		{
			`\n`, // backslash escapes are not interpreted in template literals
			nil,
			cty.StringVal("\\n"),
			0,
		},
		{
			`\uu1234`, // backslash escapes are not interpreted in template literals
			nil,       // (this is intentionally an invalid one to ensure we don't produce an error)
			cty.StringVal("\\uu1234"),
			0,
		},
		{
			`$`,
			nil,
			cty.StringVal("$"),
			0,
		},
		{
			`$$`,
			nil,
			cty.StringVal("$$"),
			0,
		},
		{
			`%`,
			nil,
			cty.StringVal("%"),
			0,
		},
		{
			`%%`,
			nil,
			cty.StringVal("%%"),
			0,
		},
		{
			`hello %%{ if true }world%%{ endif }`,
			nil,
			cty.StringVal(`hello %{ if true }world%{ endif }`),
			0,
		},
		{
			`hello $%{ if true }world%{ endif }`,
			nil,
			cty.StringVal("hello $world"),
			0,
		},
		{
			`%{ endif }`,
			nil,
			cty.UnknownVal(cty.String),
			1, // Unexpected endif directive
		},
		{
			`%{ endfor }`,
			nil,
			cty.UnknownVal(cty.String),
			1, // Unexpected endfor directive
		},
		{ // marks from uninterpolated values are ignored
			`hello%{ if false } ${target}%{ endif }`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"target": cty.StringVal("world").Mark("sensitive"),
				},
			},
			cty.StringVal("hello"),
			0,
		},
		{ // marks from interpolated values are passed through
			`${greeting} ${target}`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"greeting": cty.StringVal("hello").Mark("english"),
					"target":   cty.StringVal("world").Mark("sensitive"),
				},
			},
			cty.StringVal("hello world").WithMarks(cty.NewValueMarks("english", "sensitive")),
			0,
		},
		{ // can use marks by traversing complex values
			`Authenticate with "${secrets.passphrase}"`,
			&hcl.EvalContext{
				Variables: map[string]cty.Value{
					"secrets": cty.MapVal(map[string]cty.Value{
						"passphrase": cty.StringVal("my voice is my passport").Mark("sensitive"),
					}).Mark("sensitive"),
				},
			},
			cty.StringVal(`Authenticate with "my voice is my passport"`).WithMarks(cty.NewValueMarks("sensitive")),
			0,
		},
	}

	for _, test := range tests {
		t.Run(test.input, func(t *testing.T) {
			expr, parseDiags := ParseTemplate([]byte(test.input), "", hcl.Pos{Line: 1, Column: 1, Byte: 0})

			got, valDiags := expr.Value(test.ctx)

			diagCount := len(parseDiags) + len(valDiags)

			if diagCount != test.diagCount {
				t.Errorf("wrong number of diagnostics %d; want %d", diagCount, test.diagCount)
				for _, diag := range parseDiags {
					t.Logf(" - %s", diag.Error())
				}
				for _, diag := range valDiags {
					t.Logf(" - %s", diag.Error())
				}
			}

			if !got.RawEquals(test.want) {
				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.want)
			}
		})
	}

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"config"
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"mstree"
	"os"
	"runtime"
	"runtime/debug"
	"web"
)

func main() {
	var format string
	log := logging.MustGetLogger("metricsearch")

	var confFile, reindexFile string
	var stdinImport bool
	flag.StringVar(&confFile, "c", "/etc/metricsearch.conf", "metricsearch config filename")
	flag.StringVar(&reindexFile, "reindex", "", "reindex from plain text metrics file")
	flag.BoolVar(&stdinImport, "stdin", false, "reindex from stdin")
	flag.Parse()

	conf := config.Load(confFile)
	logging.SetLevel(conf.LogLevel, "metricsearch")

	switch conf.Log {
	case "syslog":
		backend, err := logging.NewSyslogBackend("metricsearch")
		if err != nil {
			log.Error(err.Error())
			return
		}
		logging.SetBackend(backend)
		format = "%{program} %{level} %{message}"
	case "":
		format = "%{color}%{level} %{color:reset}%{message}"
	default:
		f, err := os.OpenFile(conf.Log, os.O_APPEND|os.O_CREATE|os.O_WRONLY, os.FileMode(0644))
		if err != nil {
			log.Error(err.Error())
			return
		}
		backend := logging.NewLogBackend(f, "", 0)
		logging.SetBackend(backend)
		defer f.Close()
		format = "[%{time:2006-01-02 15:04:05}] %{level} %{message}"
	}
	logging.SetFormatter(logging.MustStringFormatter(format))

	tree, err := mstree.NewTree(conf.IndexDirectory, conf.SyncBufferSize)
	if err != nil {
		log.Critical("No way to continue, exiting.")
		return
	}

	log.Debug("Configuring runtime: GCPercent(%d), MaxCores(%d), MaxThreads(%d)", conf.GCPercent, conf.MaxCores, conf.MaxThreads)
	runtime.GOMAXPROCS(conf.MaxCores)
	debug.SetGCPercent(conf.GCPercent)
	debug.SetMaxThreads(conf.MaxThreads)

	if stdinImport {
		err := tree.DropIndex()
		if err != nil {
			log.Critical("Error dropping index: %s", err.Error())
			return
		}
		sc := bufio.NewScanner(os.Stdin)
		for sc.Scan() {
			tree.Add(sc.Text())
		}
		log.Notice("Reindexing complete")
		return
	}

	if reindexFile != "" {
		err := tree.DropIndex()
		if err != nil {
			log.Critical("Error dropping index: %s", err.Error())
			return
		}
		err = tree.LoadTxt(reindexFile, -1)
		if err != nil {
			log.Critical("Reindexing error, exiting.")
			return
		} else {
			log.Notice("Reindexing complete")
			return
		}
	} else {
		tree.LoadIndex()
		server := web.NewServer(tree)
		addr := fmt.Sprintf("%s:%d", conf.Host, conf.Port)
		server.Start(addr)
	}

}
<commit_msg>log reopen by sighup<commit_after>package main

import (
	"bufio"
	"config"
	"flag"
	"fmt"
	logging "github.com/op/go-logging"
	"mstree"
	"os"
	"os/signal"
	"runtime"
	"runtime/debug"
	"syscall"
	"web"
)

const (
	DEFAULT_CONFIG_FILE = "/etc/metricsearch.conf"
)

var (
	logfileName string
	logfile     *os.File
	log         *logging.Logger = logging.MustGetLogger("metricsearch")
)

func reopenLog() {
	var err error
	if logfile != nil {
		logfile.Close()
	}
	logfile, err = os.OpenFile(logfileName, os.O_APPEND|os.O_CREATE|os.O_WRONLY, os.FileMode(0644))
	if err != nil {
		return
	}
	backend := logging.NewLogBackend(logfile, "", 0)
	logging.SetBackend(backend)
}

func hupCatcher() {
	c := make(chan os.Signal)
	signal.Notify(c, syscall.SIGHUP)
	for _ = range c {
		log.Debug("HUP signal catched, reopening logfile %s", logfileName)
		reopenLog()
	}
}

func main() {
	var format string
	var confFile, reindexFile string
	var stdinImport bool
	flag.StringVar(&confFile, "c", DEFAULT_CONFIG_FILE, "metricsearch config filename")
	flag.StringVar(&reindexFile, "reindex", "", "reindex from plain text metrics file")
	flag.BoolVar(&stdinImport, "stdin", false, "reindex from stdin")
	flag.Parse()

	conf := config.Load(confFile)
	logging.SetLevel(conf.LogLevel, "metricsearch")

	switch conf.Log {
	case "syslog":
		backend, err := logging.NewSyslogBackend("metricsearch")
		if err != nil {
			log.Error(err.Error())
			return
		}
		logging.SetBackend(backend)
		format = "%{program} %{level} %{message}"
	case "":
		format = "%{color}%{level} %{color:reset}%{message}"
	default:
		logfileName = conf.Log
		go hupCatcher()
		format = "[%{time:2006-01-02 15:04:05}] %{level} %{message}"
	}
	logging.SetFormatter(logging.MustStringFormatter(format))

	tree, err := mstree.NewTree(conf.IndexDirectory, conf.SyncBufferSize)
	if err != nil {
		log.Critical("No way to continue, exiting.")
		return
	}

	log.Debug("Configuring runtime: GCPercent(%d), MaxCores(%d), MaxThreads(%d)", conf.GCPercent, conf.MaxCores, conf.MaxThreads)
	runtime.GOMAXPROCS(conf.MaxCores)
	debug.SetGCPercent(conf.GCPercent)
	debug.SetMaxThreads(conf.MaxThreads)

	if stdinImport {
		err := tree.DropIndex()
		if err != nil {
			log.Critical("Error dropping index: %s", err.Error())
			return
		}
		sc := bufio.NewScanner(os.Stdin)
		for sc.Scan() {
			tree.Add(sc.Text())
		}
		log.Notice("Reindexing complete")
		return
	}

	if reindexFile != "" {
		err := tree.DropIndex()
		if err != nil {
			log.Critical("Error dropping index: %s", err.Error())
			return
		}
		err = tree.LoadTxt(reindexFile, -1)
		if err != nil {
			log.Critical("Reindexing error, exiting.")
			return
		} else {
			log.Notice("Reindexing complete")
			return
		}
	} else {
		tree.LoadIndex()
		server := web.NewServer(tree)
		addr := fmt.Sprintf("%s:%d", conf.Host, conf.Port)
		server.Start(addr)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"github.com/gopherjs/gopherjs/js"
	"github.com/jaracil/nxcli"
	"github.com/jaracil/nxcli/nxcore"
)

func main() {

	// Node.js
	if require := js.Global.Get("require"); require != js.Undefined {
		if ws := require.Invoke("ws"); ws != js.Undefined {
			println("Running on Node.js with websocket module: 'ws'")

			if removeListener := ws.Get("prototype").Get("removeEventListener"); removeListener == js.Undefined {
				println("Missing removeEventListener! Patching WebSocket library...")

				ws.Get("prototype").Set("removeEventListener", js.MakeFunc(func(this *js.Object, args []*js.Object) interface{} {
					this.Call("removeListener", args[0], args[1])
					return nil
				}))

				println("OK")
			}

			js.Global.Set("WebSocket", ws)
		} else {
			println("Running on Node.js but websocket module missing")
		}
	}

	dial := func(a string, cb func(interface{}, error)) {
		go func() {
			nc, e := nexus.Dial(a, nil)
			go cb(WrapNexusConn(nc), e)
		}()
	}

	var nexus *js.Object

	if js.Module == js.Undefined {
		// Browser
		js.Global.Set("nexus", make(map[interface{}]interface{}))
		nexus = js.Global.Get("nexus")
	} else {
		// Node.js
		nexus = js.Module.Get("exports")
	}

	nexus.Set("dial", dial)
	nexus.Set("errors", nxcore.ErrStr)

	println("Nexus Client loaded")
}
<commit_msg>Removed extra goroutine call in Dial method<commit_after>package main

import (
	"github.com/gopherjs/gopherjs/js"
	"github.com/jaracil/nxcli"
	"github.com/jaracil/nxcli/nxcore"
)

func main() {

	// Node.js
	if require := js.Global.Get("require"); require != js.Undefined {
		if ws := require.Invoke("ws"); ws != js.Undefined {
			println("Running on Node.js with websocket module: 'ws'")

			if removeListener := ws.Get("prototype").Get("removeEventListener"); removeListener == js.Undefined {
				println("Missing removeEventListener! Patching WebSocket library...")

				ws.Get("prototype").Set("removeEventListener", js.MakeFunc(func(this *js.Object, args []*js.Object) interface{} {
					this.Call("removeListener", args[0], args[1])
					return nil
				}))

				println("OK")
			}

			js.Global.Set("WebSocket", ws)
		} else {
			println("Running on Node.js but websocket module missing")
		}
	}

	dial := func(a string, cb func(interface{}, error)) {
		go func() {
			nc, e := nexus.Dial(a, nil)
			cb(WrapNexusConn(nc), e)
		}()
	}

	var nexus *js.Object

	if js.Module == js.Undefined {
		// Browser
		js.Global.Set("nexus", make(map[interface{}]interface{}))
		nexus = js.Global.Get("nexus")
	} else {
		// Node.js
		nexus = js.Module.Get("exports")
	}

	nexus.Set("dial", dial)
	nexus.Set("errors", nxcore.ErrStr)

	println("Nexus Client loaded")
}
<|endoftext|>


<commit_before>package cluster

import (
	"testing"

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

func TestBuildContainerConfig(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Equal(t, len(config.Env), 0)
	assert.Equal(t, len(config.Labels), 0)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true"}})
	assert.Equal(t, len(config.Env), 1)
	assert.Equal(t, len(config.Labels), 0)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"constraint:test==true"}})
	assert.Equal(t, len(config.Env), 0)
	assert.Equal(t, len(config.Labels), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:container==test"}})
	assert.Equal(t, len(config.Env), 0)
	assert.Equal(t, len(config.Labels), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Equal(t, len(config.Env), 1)
	assert.Equal(t, len(config.Labels), 2)
}

func TestConstraints(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Equal(t, len(config.Constraints()), 0)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"constraint:test==true"}})
	assert.Equal(t, len(config.Constraints()), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Equal(t, len(config.Constraints()), 1)
}

func TestAffinities(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Equal(t, len(config.Affinities()), 0)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:container==test"}})
	assert.Equal(t, len(config.Affinities()), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Equal(t, len(config.Affinities()), 1)
}
<commit_msg>cleanup config asserts<commit_after>package cluster

import (
	"testing"

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

func TestBuildContainerConfig(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Empty(t, config.Env)
	assert.Empty(t, config.Labels)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true"}})
	assert.Len(t, config.Env, 1)
	assert.Empty(t, config.Labels)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"constraint:test==true"}})
	assert.Empty(t, config.Env)
	assert.Len(t, config.Labels, 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:container==test"}})
	assert.Empty(t, config.Env)
	assert.Len(t, config.Labels, 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Len(t, config.Env, 1)
	assert.Len(t, config.Labels, 2)
}

func TestConstraints(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Empty(t, config.Constraints())

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"constraint:test==true"}})
	assert.Len(t, config.Constraints(), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Len(t, config.Constraints(), 1)
}

func TestAffinities(t *testing.T) {
	config := BuildContainerConfig(&dockerclient.ContainerConfig{})
	assert.Empty(t, config.Affinities())

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"affinity:container==test"}})
	assert.Len(t, config.Affinities(), 1)

	config = BuildContainerConfig(&dockerclient.ContainerConfig{Env: []string{"test=true", "constraint:test==true", "affinity:container==test"}})
	assert.Len(t, config.Affinities(), 1)
}
<|endoftext|>


<commit_before>package scrape

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"io"
	"io/ioutil"
	"strings"

	"github.com/clbanning/mxj"
)

func newEncoder(s Task) (e encoder) {
	switch strings.ToLower(s.Scrapers[0].Opts.Format) {
	case "csv":
		e = CSVEncoder{
			comma:     ",",
			partNames: s.Scrapers[0].partNames(),
			results:   &s.Results,
		}
		return
	case "json":
		e = JSONEncoder{
			paginateResults: s.Scrapers[0].Opts.PaginateResults,
			results:         &s.Results,
		}
		return
	case "xml":
		e = XMLEncoder{
			results: &s.Results,
		}
		return e
	default:
		return nil
	}
	return nil
}

type encoder interface {
	Encode() (io.ReadCloser, error)
}

// CSVEncoder transforms parsed data to CSV format.
type CSVEncoder struct {
	partNames []string
	comma     string
	results   *Results
}

// JSONEncoder transforms parsed data to JSON format.
type JSONEncoder struct {
	paginateResults bool
	results         *Results
}

// XMLEncoder transforms parsed data to XML format.
type XMLEncoder struct {
	results *Results
}

//Encode method implementation for JSONEncoder
func (e JSONEncoder) Encode() (io.ReadCloser, error) {
	var buf bytes.Buffer
	if e.paginateResults {
		json.NewEncoder(&buf).Encode(e.results)
	} else {
		json.NewEncoder(&buf).Encode(e.results.AllBlocks())
	}
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//Encode method implementation for CSVEncoder
func (e CSVEncoder) Encode() (io.ReadCloser, error) {
	var buf bytes.Buffer
	/*
		includeHeader := true
		w := csv.NewWriter(&buf)
		for i, page := range results.Results {
			if i != 0 {
				includeHeader = false
			}
			err = encodeCSV(names, includeHeader, page, ",", w)
			if err != nil {
				logger.Error(err)
			}
		}
		w.Flush()
	*/
	w := csv.NewWriter(&buf)

	err := encodeCSV(e.partNames, e.results.AllBlocks(), e.comma, w)
	if err != nil {
		return nil, err
	}
	w.Flush()
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//Encode method implementation for XMLEncoder
func (e XMLEncoder) Encode() (io.ReadCloser, error) {
	/*
		case "xmlviajson":
			var jbuf bytes.Buffer
			if config.Opts.PaginateResults {
				json.NewEncoder(&jbuf).Encode(results)
			} else {
				json.NewEncoder(&jbuf).Encode(results.AllBlocks())
			}
			//var buf bytes.Buffer
			m, err := mxj.NewMapJson(jbuf.Bytes())
			err = m.XmlIndentWriter(&buf, "", "  ")
			if err != nil {
				logger.Error(err)
			}
	*/
	var buf bytes.Buffer
	err := encodeXML(e.results.AllBlocks(), &buf)
	if err != nil {
		return nil, err
	}
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//encodeCSV writes data to w *csv.Writer.
//header represent an array of fields for csv.
//rows store csv records to be written.
//comma is a separator between record fields. Default value is ","
func encodeCSV(header []string, rows []map[string]interface{}, comma string, w *csv.Writer) error {
	if comma == "" {
		comma = ","
	}
	w.Comma = rune(comma[0])
	//Add Header string to csv or no
	if len(header) > 0 {
		if err := w.Write(header); err != nil {
			return err
		}
	}
	r := make([]string, len(header))
	for _, row := range rows {
		for i, column := range header {
			switch v := row[column].(type) {
			case string:
				r[i] = v
			case []string:
				r[i] = strings.Join(v, ";")
			case nil:
				r[i] = ""
			}
		}
		if err := w.Write(r); err != nil {
			return err
		}
	}
	return nil
}

//encodeXML writes data blocks to XML file.
func encodeXML(blocks []map[string]interface{}, buf *bytes.Buffer) error {
	mxj.XMLEscapeChars(true)
	//write header to xml
	buf.Write([]byte(`<?xml version="1.0" encoding="UTF-8"?>`))
	buf.Write([]byte("<items>"))
	for _, elem := range blocks {
		m := mxj.Map(elem)
		//err := m.XmlIndentWriter(&buf, "", "  ", "object")
		err := m.XmlWriter(buf, "item")
		if err != nil {
			return err
		}
	}
	buf.Write([]byte("</items>"))
	return nil
}
<commit_msg>added Format method to encoders<commit_after>package scrape

import (
	"bytes"
	"encoding/csv"
	"encoding/json"
	"io"
	"io/ioutil"
	"strings"

	"github.com/clbanning/mxj"
)

func newEncoder(s Task) (e encoder) {
	switch strings.ToLower(s.Scrapers[0].Opts.Format) {
	case "csv":
		e = CSVEncoder{
			comma:     ",",
			partNames: s.Scrapers[0].partNames(),
			results:   &s.Results,
		}
		return
	case "json":
		e = JSONEncoder{
			paginateResults: s.Scrapers[0].Opts.PaginateResults,
			results:         &s.Results,
		}
		return
	case "xml":
		e = XMLEncoder{
			results: &s.Results,
		}
		return e
	default:
		return nil
	}
	return nil
}

type encoder interface {
	Encode() (io.ReadCloser, error)
	Format() string
}

// CSVEncoder transforms parsed data to CSV format.
type CSVEncoder struct {
	partNames []string
	comma     string
	results   *Results
}

// JSONEncoder transforms parsed data to JSON format.
type JSONEncoder struct {
	paginateResults bool
	results         *Results
}

// XMLEncoder transforms parsed data to XML format.
type XMLEncoder struct {
	results *Results
}

//Encode method implementation for JSONEncoder
func (e JSONEncoder) Encode() (io.ReadCloser, error) {
	var buf bytes.Buffer
	if e.paginateResults {
		json.NewEncoder(&buf).Encode(e.results)
	} else {
		json.NewEncoder(&buf).Encode(e.results.AllBlocks())
	}
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//Format returns format of JSONEncoder
func (JSONEncoder) Format() string {
	return "JSON"
}

//Encode method implementation for CSVEncoder
func (e CSVEncoder) Encode() (io.ReadCloser, error) {
	var buf bytes.Buffer
	/*
		includeHeader := true
		w := csv.NewWriter(&buf)
		for i, page := range results.Results {
			if i != 0 {
				includeHeader = false
			}
			err = encodeCSV(names, includeHeader, page, ",", w)
			if err != nil {
				logger.Error(err)
			}
		}
		w.Flush()
	*/
	w := csv.NewWriter(&buf)

	err := encodeCSV(e.partNames, e.results.AllBlocks(), e.comma, w)
	if err != nil {
		return nil, err
	}
	w.Flush()
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//Format returns format of CSVEncoder
func (CSVEncoder) Format() string {
	return "CSV"
}

//Encode method implementation for XMLEncoder
func (e XMLEncoder) Encode() (io.ReadCloser, error) {
	/*
		case "xmlviajson":
			var jbuf bytes.Buffer
			if config.Opts.PaginateResults {
				json.NewEncoder(&jbuf).Encode(results)
			} else {
				json.NewEncoder(&jbuf).Encode(results.AllBlocks())
			}
			//var buf bytes.Buffer
			m, err := mxj.NewMapJson(jbuf.Bytes())
			err = m.XmlIndentWriter(&buf, "", "  ")
			if err != nil {
				logger.Error(err)
			}
	*/
	var buf bytes.Buffer
	err := encodeXML(e.results.AllBlocks(), &buf)
	if err != nil {
		return nil, err
	}
	readCloser := ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
	return readCloser, nil
}

//Format returns format of XMLEncoder
func (XMLEncoder) Format() string {
	return "XML"
}

//encodeCSV writes data to w *csv.Writer.
//header represent an array of fields for csv.
//rows store csv records to be written.
//comma is a separator between record fields. Default value is ","
func encodeCSV(header []string, rows []map[string]interface{}, comma string, w *csv.Writer) error {
	if comma == "" {
		comma = ","
	}
	w.Comma = rune(comma[0])
	//Add Header string to csv or no
	if len(header) > 0 {
		if err := w.Write(header); err != nil {
			return err
		}
	}
	r := make([]string, len(header))
	for _, row := range rows {
		for i, column := range header {
			switch v := row[column].(type) {
			case string:
				r[i] = v
			case []string:
				r[i] = strings.Join(v, ";")
			case nil:
				r[i] = ""
			}
		}
		if err := w.Write(r); err != nil {
			return err
		}
	}
	return nil
}

//encodeXML writes data blocks to XML file.
func encodeXML(blocks []map[string]interface{}, buf *bytes.Buffer) error {
	mxj.XMLEscapeChars(true)
	//write header to xml
	buf.Write([]byte(`<?xml version="1.0" encoding="UTF-8"?>`))
	buf.Write([]byte("<items>"))
	for _, elem := range blocks {
		m := mxj.Map(elem)
		//err := m.XmlIndentWriter(&buf, "", "  ", "object")
		err := m.XmlWriter(buf, "item")
		if err != nil {
			return err
		}
	}
	buf.Write([]byte("</items>"))
	return nil
}
<|endoftext|>


<commit_before>package next

type TreeNode struct {
	Val               int
	Left, Right, Next *TreeNode
}

func connect(root *TreeNode) *TreeNode {
	return root
}
<commit_msg>solve 116 use levelorder traversal not constant space<commit_after>package next

type TreeNode struct {
	Val               int
	Left, Right, Next *TreeNode
}

func connect(root *TreeNode) *TreeNode {
	return uselevelorder(root)
}

// uselevelorder time complexity O(N), space complexity O(N) which is not constant.
func uselevelorder(root *TreeNode) *TreeNode {
	if root == nil || root.Left == nil && root.Right == nil {
		return root
	}
	var st []*TreeNode
	st = append(st, root)
	nst := len(st)
	for nst != 0 {
		var tmp []*TreeNode
		for i, node := range st {
			if node.Left != nil && node.Right != nil {
				tmp = append(tmp, node.Left, node.Right)
			}
			if i+1 < nst {
				node.Next = st[i+1]
			}
		}
		st = tmp
		nst = len(st)
	}
	return root
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"github.com/timeredbull/tsuru/cmd"
	. "launchpad.net/gocheck"
	"net/http"
)

func (s *S) TestAddKey(c *C) {
	expected := "Key successfully added!\n"
	context := cmd.Context{
		Cmds:   []string{},
		Args:   []string{},
		Stdout: manager.Stdout,
		Stderr: manager.Stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "", status: http.StatusOK}})
	command := AddKeyCommand{}
	err := command.Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(manager.Stdout.(*bytes.Buffer).String(), Equals, expected)
}

func (s *S) TestRemoveKey(c *C) {
	expected := "Key successfully removed!\n"
	context := cmd.Context{
		Cmds:   []string{},
		Args:   []string{},
		Stdout: manager.Stdout,
		Stderr: manager.Stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "", status: http.StatusOK}})
	command := RemoveKey{}
	err := command.Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(manager.Stdout.(*bytes.Buffer).String(), Equals, expected)
}

func (s *S) TestKey(c *C) {
	expect := map[string]interface{}{
		"add":    &AddKeyCommand{},
		"remove": &RemoveKey{},
	}
	command := Key{}
	c.Assert(command.Subcommands(), DeepEquals, expect)
}
<commit_msg>cmd/appcli: added message<commit_after>package main

import (
	"bytes"
	"github.com/timeredbull/tsuru/cmd"
	. "launchpad.net/gocheck"
	"net/http"
)

func (s *S) TestAddKey(c *C) {
	expected := "Key successfully added!\n"
	context := cmd.Context{
		Cmds:   []string{},
		Args:   []string{},
		Stdout: manager.Stdout,
		Stderr: manager.Stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "success", status: http.StatusOK}})
	command := AddKeyCommand{}
	err := command.Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(manager.Stdout.(*bytes.Buffer).String(), Equals, expected)
}

func (s *S) TestRemoveKey(c *C) {
	expected := "Key successfully removed!\n"
	context := cmd.Context{
		Cmds:   []string{},
		Args:   []string{},
		Stdout: manager.Stdout,
		Stderr: manager.Stderr,
	}
	client := cmd.NewClient(&http.Client{Transport: &transport{msg: "success", status: http.StatusOK}})
	command := RemoveKey{}
	err := command.Run(&context, client)
	c.Assert(err, IsNil)
	c.Assert(manager.Stdout.(*bytes.Buffer).String(), Equals, expected)
}

func (s *S) TestKey(c *C) {
	expect := map[string]interface{}{
		"add":    &AddKeyCommand{},
		"remove": &RemoveKey{},
	}
	command := Key{}
	c.Assert(command.Subcommands(), DeepEquals, expect)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/nu7hatch/gouuid"

	"github.com/cloudfoundry-incubator/auctioneer/auctionmetricemitterdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/handlers"
	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/locket/presence"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"

	"github.com/cloudfoundry-incubator/auction/auctionrunner"
	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gunk/workpool"
	"github.com/pivotal-golang/clock"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server addresses (ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:9016",
	"host:port to serve auction and LRP stop requests on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

const (
	auctionRunnerTimeout      = 10 * time.Second
	auctionRunnerWorkPoolSize = 1000
	dropsondeDestination      = "localhost:3457"
	dropsondeOrigin           = "auctioneer"
	serverProtocol            = "http"
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("auctioneer")
	initializeDropsonde(logger)

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSession("auctioneer", *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	locketClient := locket.NewClient(consulSession, clock.NewClock(), logger)
	bbsClient := bbs.NewClient(*bbsAddress)

	auctionRunner := initializeAuctionRunner(bbsClient, locketClient, consulSession, logger)
	auctionServer := initializeAuctionServer(auctionRunner, logger)
	lockMaintainer := initializeLockMaintainer(locketClient, logger)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"auction-runner", auctionRunner},
		{"auction-server", auctionServer},
	}

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeAuctionRunner(bbsClient bbs.Client, locketClient locket.Client, consulSession *consuladapter.Session, logger lager.Logger) auctiontypes.AuctionRunner {
	httpClient := cf_http.NewClient()
	repClientFactory := rep.NewClientFactory(httpClient)

	delegate := auctionrunnerdelegate.New(repClientFactory, bbsClient, locketClient, logger)
	metricEmitter := auctionmetricemitterdelegate.New()
	workPool, err := workpool.NewWorkPool(auctionRunnerWorkPoolSize)
	if err != nil {
		logger.Fatal("failed-to-construct-auction-runner-workpool", err, lager.Data{"num-workers": auctionRunnerWorkPoolSize}) // should never happen
	}

	return auctionrunner.New(
		delegate,
		metricEmitter,
		clock.NewClock(),
		workPool,
		logger,
	)
}

func initializeDropsonde(logger lager.Logger) {
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeAuctionServer(runner auctiontypes.AuctionRunner, logger lager.Logger) ifrit.Runner {
	return http_server.New(*listenAddr, handlers.New(runner, logger))
}

func initializeLockMaintainer(locket locket.Client, logger lager.Logger) ifrit.Runner {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}

	localIP, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("Couldn't determine local IP", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("%s://%s:%s", serverProtocol, localIP, port)

	auctioneerPresence := presence.NewAuctioneerPresence(uuid.String(), address)
	lockMaintainer, err := locket.NewAuctioneerLock(auctioneerPresence, *lockRetryInterval)
	if err != nil {
		logger.Fatal("Couldn't create lock maintainer", err)
	}

	return lockMaintainer
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}
<commit_msg>support mutual-auth TLS to BBS<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"
	"time"

	"github.com/nu7hatch/gouuid"

	"github.com/cloudfoundry-incubator/auctioneer/auctionmetricemitterdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/auctionrunnerdelegate"
	"github.com/cloudfoundry-incubator/auctioneer/handlers"
	"github.com/cloudfoundry-incubator/bbs"
	"github.com/cloudfoundry-incubator/cf-debug-server"
	cf_lager "github.com/cloudfoundry-incubator/cf-lager"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/cloudfoundry-incubator/consuladapter"
	"github.com/cloudfoundry-incubator/locket"
	"github.com/cloudfoundry-incubator/locket/presence"
	"github.com/cloudfoundry-incubator/rep"
	"github.com/pivotal-golang/lager"
	"github.com/pivotal-golang/localip"

	"github.com/cloudfoundry-incubator/auction/auctionrunner"
	"github.com/cloudfoundry-incubator/auction/auctiontypes"
	"github.com/cloudfoundry/dropsonde"
	"github.com/cloudfoundry/gunk/workpool"
	"github.com/pivotal-golang/clock"
	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/grouper"
	"github.com/tedsuo/ifrit/http_server"
	"github.com/tedsuo/ifrit/sigmon"
)

var communicationTimeout = flag.Duration(
	"communicationTimeout",
	10*time.Second,
	"Timeout applied to all HTTP requests.",
)

var consulCluster = flag.String(
	"consulCluster",
	"",
	"comma-separated list of consul server addresses (ip:port)",
)

var lockTTL = flag.Duration(
	"lockTTL",
	locket.LockTTL,
	"TTL for service lock",
)

var lockRetryInterval = flag.Duration(
	"lockRetryInterval",
	locket.RetryInterval,
	"interval to wait before retrying a failed lock acquisition",
)

var listenAddr = flag.String(
	"listenAddr",
	"0.0.0.0:9016",
	"host:port to serve auction and LRP stop requests on",
)

var bbsAddress = flag.String(
	"bbsAddress",
	"",
	"Address to the BBS Server",
)

var bbsCACert = flag.String(
	"bbsCACert",
	"",
	"path to certificate authority cert used for mutually authenticated TLS BBS communication",
)

var bbsClientCert = flag.String(
	"bbsClientCert",
	"",
	"path to client cert used for mutually authenticated TLS BBS communication",
)

var bbsClientKey = flag.String(
	"bbsClientKey",
	"",
	"path to client key used for mutually authenticated TLS BBS communication",
)

const (
	auctionRunnerTimeout      = 10 * time.Second
	auctionRunnerWorkPoolSize = 1000
	dropsondeDestination      = "localhost:3457"
	dropsondeOrigin           = "auctioneer"
	serverProtocol            = "http"
)

func main() {
	cf_debug_server.AddFlags(flag.CommandLine)
	cf_lager.AddFlags(flag.CommandLine)
	flag.Parse()

	cf_http.Initialize(*communicationTimeout)

	logger, reconfigurableSink := cf_lager.New("auctioneer")
	initializeDropsonde(logger)

	if err := validateBBSAddress(); err != nil {
		logger.Fatal("invalid-bbs-address", err)
	}

	client, err := consuladapter.NewClient(*consulCluster)
	if err != nil {
		logger.Fatal("new-client-failed", err)
	}

	sessionMgr := consuladapter.NewSessionManager(client)
	consulSession, err := consuladapter.NewSession("auctioneer", *lockTTL, client, sessionMgr)
	if err != nil {
		logger.Fatal("consul-session-failed", err)
	}

	locketClient := locket.NewClient(consulSession, clock.NewClock(), logger)

	auctionRunner := initializeAuctionRunner(initializeBBSClient(logger), locketClient, consulSession, logger)
	auctionServer := initializeAuctionServer(auctionRunner, logger)
	lockMaintainer := initializeLockMaintainer(locketClient, logger)

	members := grouper.Members{
		{"lock-maintainer", lockMaintainer},
		{"auction-runner", auctionRunner},
		{"auction-server", auctionServer},
	}

	if dbgAddr := cf_debug_server.DebugAddress(flag.CommandLine); dbgAddr != "" {
		members = append(grouper.Members{
			{"debug-server", cf_debug_server.Runner(dbgAddr, reconfigurableSink)},
		}, members...)
	}

	group := grouper.NewOrdered(os.Interrupt, members)

	monitor := ifrit.Invoke(sigmon.New(group))

	logger.Info("started")

	err = <-monitor.Wait()
	if err != nil {
		logger.Error("exited-with-failure", err)
		os.Exit(1)
	}

	logger.Info("exited")
}

func initializeAuctionRunner(bbsClient bbs.Client, locketClient locket.Client, consulSession *consuladapter.Session, logger lager.Logger) auctiontypes.AuctionRunner {
	httpClient := cf_http.NewClient()
	repClientFactory := rep.NewClientFactory(httpClient)

	delegate := auctionrunnerdelegate.New(repClientFactory, bbsClient, locketClient, logger)
	metricEmitter := auctionmetricemitterdelegate.New()
	workPool, err := workpool.NewWorkPool(auctionRunnerWorkPoolSize)
	if err != nil {
		logger.Fatal("failed-to-construct-auction-runner-workpool", err, lager.Data{"num-workers": auctionRunnerWorkPoolSize}) // should never happen
	}

	return auctionrunner.New(
		delegate,
		metricEmitter,
		clock.NewClock(),
		workPool,
		logger,
	)
}

func initializeDropsonde(logger lager.Logger) {
	err := dropsonde.Initialize(dropsondeDestination, dropsondeOrigin)
	if err != nil {
		logger.Error("failed to initialize dropsonde: %v", err)
	}
}

func initializeAuctionServer(runner auctiontypes.AuctionRunner, logger lager.Logger) ifrit.Runner {
	return http_server.New(*listenAddr, handlers.New(runner, logger))
}

func initializeLockMaintainer(locket locket.Client, logger lager.Logger) ifrit.Runner {
	uuid, err := uuid.NewV4()
	if err != nil {
		logger.Fatal("Couldn't generate uuid", err)
	}

	localIP, err := localip.LocalIP()
	if err != nil {
		logger.Fatal("Couldn't determine local IP", err)
	}

	port := strings.Split(*listenAddr, ":")[1]
	address := fmt.Sprintf("%s://%s:%s", serverProtocol, localIP, port)

	auctioneerPresence := presence.NewAuctioneerPresence(uuid.String(), address)
	lockMaintainer, err := locket.NewAuctioneerLock(auctioneerPresence, *lockRetryInterval)
	if err != nil {
		logger.Fatal("Couldn't create lock maintainer", err)
	}

	return lockMaintainer
}

func validateBBSAddress() error {
	if *bbsAddress == "" {
		return errors.New("bbsAddress is required")
	}
	return nil
}

func initializeBBSClient(logger lager.Logger) bbs.Client {
	bbsURL, err := url.Parse(*bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(*bbsAddress)
	}

	bbsClient, err := bbs.NewSecureClient(*bbsAddress, *bbsCACert, *bbsClientCert, *bbsClientKey)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}
<|endoftext|>


<commit_before>/*

	Autoreader is a simple program, designed to be run from go:generate, that
	helps generate the annoying boilerplate to implement
	boardgame.PropertyReader and boardgame.PropertyReadSetter.

*/
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/MarcGrol/golangAnnotations/parser"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"
)

var headerTemplate *template.Template
var readerTemplate *template.Template
var readSetterTemplate *template.Template

const magicDocLinePrefix = "+autoreader"

type appOptions struct {
	OutputFile       string
	PackageDirectory string
	Help             bool
	flagSet          *flag.FlagSet
}

type templateConfig struct {
	FirstLetter string
	StructName  string
}

func init() {
	headerTemplate = template.Must(template.New("header").Parse(headerTemplateText))
	readerTemplate = template.Must(template.New("reader").Parse(readerTemplateText))
	readSetterTemplate = template.Must(template.New("readsetter").Parse(readSetterTemplateText))
}

func defineFlags(options *appOptions) {
	options.flagSet.StringVar(&options.OutputFile, "out", "auto_reader.go", "Defines which file to render output to. WARNING: it will be overwritten!")
	options.flagSet.StringVar(&options.PackageDirectory, "pkg", "examplepkg/", "Which package to process")
	options.flagSet.BoolVar(&options.Help, "h", false, "If set, print help message and quit.")
}

func getOptions(flagSet *flag.FlagSet, flagArguments []string) *appOptions {
	options := &appOptions{flagSet: flagSet}
	defineFlags(options)
	flagSet.Parse(flagArguments)
	return options
}

func main() {
	flagSet := flag.CommandLine
	process(getOptions(flagSet, os.Args[1:]), os.Stderr)
}

func process(options *appOptions, errOut io.ReadWriter) {

	if options.Help {
		options.flagSet.SetOutput(errOut)
		options.flagSet.PrintDefaults()
		return
	}

	output, err := processPackage(options.PackageDirectory)

	if err != nil {
		fmt.Fprintln(errOut, "ERROR", err)
		return
	}

	ioutil.WriteFile(options.OutputFile, []byte(output), 0644)
}

func processPackage(location string) (output string, err error) {
	sources, err := parser.ParseSourceDir(location, ".*")

	if err != nil {
		return "", errors.New("Couldn't parse sources: " + err.Error())
	}

	haveOutputHeader := false

	for _, theStruct := range sources.Structs {

		if !haveOutputHeader {
			output += headerForPackage(theStruct.PackageName)
			haveOutputHeader = true
		}

		enableAutoReader := false

		for _, docLine := range theStruct.DocLines {
			docLine = strings.TrimPrefix(docLine, "//")
			docLine = strings.TrimSpace(docLine)
			if strings.HasPrefix(docLine, magicDocLinePrefix) {
				enableAutoReader = true
				break
			}
		}

		if !enableAutoReader {
			continue
		}

		output += readerForStruct(theStruct.Name)
		output += readSetterForStruct(theStruct.Name)
	}

	return output, nil
}

func templateOutput(template *template.Template, values interface{}) string {
	buf := new(bytes.Buffer)

	err := template.Execute(buf, values)

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

	return buf.String()
}

func headerForPackage(packageName string) string {
	return templateOutput(headerTemplate, map[string]string{
		"packageName": packageName,
	})
}

func readerForStruct(structName string) string {

	return templateOutput(readerTemplate, templateConfig{
		FirstLetter: structName[:1],
		StructName:  structName,
	})

}

func readSetterForStruct(structName string) string {
	return templateOutput(readSetterTemplate, templateConfig{
		FirstLetter: structName[:1],
		StructName:  structName,
	})
}

const headerTemplateText = `/************************************
 *
 * This file was auto-generated by autoreader. DO NOT EDIT.
 *
 ************************************/
package {{.packageName}}

import (
	"github.com/jkomoros/boardgame"
)

`

const readerTemplateText = `func ({{.FirstLetter}} *{{.StructName}}) Reader() boardgame.PropertyReader {
	return boardgame.DefaultReader({{.FirstLetter}})
}

`

const readSetterTemplateText = `func ({{.FirstLetter}} *{{.StructName}}) ReadSetter() boardgame.PropertyReadSetter {
	return boardgame.DefaultReadSetter({{.FirstLetter}})
}

`
<commit_msg>Add an option to print to console instead of printing to the file. Part of #301.<commit_after>/*

	Autoreader is a simple program, designed to be run from go:generate, that
	helps generate the annoying boilerplate to implement
	boardgame.PropertyReader and boardgame.PropertyReadSetter.

*/
package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"github.com/MarcGrol/golangAnnotations/parser"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"text/template"
)

var headerTemplate *template.Template
var readerTemplate *template.Template
var readSetterTemplate *template.Template

const magicDocLinePrefix = "+autoreader"

type appOptions struct {
	OutputFile       string
	PackageDirectory string
	PrintToConsole   bool
	Help             bool
	flagSet          *flag.FlagSet
}

type templateConfig struct {
	FirstLetter string
	StructName  string
}

func init() {
	headerTemplate = template.Must(template.New("header").Parse(headerTemplateText))
	readerTemplate = template.Must(template.New("reader").Parse(readerTemplateText))
	readSetterTemplate = template.Must(template.New("readsetter").Parse(readSetterTemplateText))
}

func defineFlags(options *appOptions) {
	options.flagSet.StringVar(&options.OutputFile, "out", "auto_reader.go", "Defines which file to render output to. WARNING: it will be overwritten!")
	options.flagSet.StringVar(&options.PackageDirectory, "pkg", "examplepkg/", "Which package to process")
	options.flagSet.BoolVar(&options.Help, "h", false, "If set, print help message and quit.")
	options.flagSet.BoolVar(&options.PrintToConsole, "print", false, "If true, will print result to console instead of writing to out.")
}

func getOptions(flagSet *flag.FlagSet, flagArguments []string) *appOptions {
	options := &appOptions{flagSet: flagSet}
	defineFlags(options)
	flagSet.Parse(flagArguments)
	return options
}

func main() {
	flagSet := flag.CommandLine
	process(getOptions(flagSet, os.Args[1:]), os.Stderr)
}

func process(options *appOptions, errOut io.ReadWriter) {

	if options.Help {
		options.flagSet.SetOutput(errOut)
		options.flagSet.PrintDefaults()
		return
	}

	output, err := processPackage(options.PackageDirectory)

	if err != nil {
		fmt.Fprintln(errOut, "ERROR", err)
		return
	}

	if options.PrintToConsole {
		fmt.Fprintln(errOut, output)
	} else {
		ioutil.WriteFile(options.OutputFile, []byte(output), 0644)
	}

}

func processPackage(location string) (output string, err error) {
	sources, err := parser.ParseSourceDir(location, ".*")

	if err != nil {
		return "", errors.New("Couldn't parse sources: " + err.Error())
	}

	haveOutputHeader := false

	for _, theStruct := range sources.Structs {

		if !haveOutputHeader {
			output += headerForPackage(theStruct.PackageName)
			haveOutputHeader = true
		}

		enableAutoReader := false

		for _, docLine := range theStruct.DocLines {
			docLine = strings.TrimPrefix(docLine, "//")
			docLine = strings.TrimSpace(docLine)
			if strings.HasPrefix(docLine, magicDocLinePrefix) {
				enableAutoReader = true
				break
			}
		}

		if !enableAutoReader {
			continue
		}

		output += readerForStruct(theStruct.Name)
		output += readSetterForStruct(theStruct.Name)
	}

	return output, nil
}

func templateOutput(template *template.Template, values interface{}) string {
	buf := new(bytes.Buffer)

	err := template.Execute(buf, values)

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

	return buf.String()
}

func headerForPackage(packageName string) string {
	return templateOutput(headerTemplate, map[string]string{
		"packageName": packageName,
	})
}

func readerForStruct(structName string) string {

	return templateOutput(readerTemplate, templateConfig{
		FirstLetter: structName[:1],
		StructName:  structName,
	})

}

func readSetterForStruct(structName string) string {
	return templateOutput(readSetterTemplate, templateConfig{
		FirstLetter: structName[:1],
		StructName:  structName,
	})
}

const headerTemplateText = `/************************************
 *
 * This file was auto-generated by autoreader. DO NOT EDIT.
 *
 ************************************/
package {{.packageName}}

import (
	"github.com/jkomoros/boardgame"
)

`

const readerTemplateText = `func ({{.FirstLetter}} *{{.StructName}}) Reader() boardgame.PropertyReader {
	return boardgame.DefaultReader({{.FirstLetter}})
}

`

const readSetterTemplateText = `func ({{.FirstLetter}} *{{.StructName}}) ReadSetter() boardgame.PropertyReadSetter {
	return boardgame.DefaultReadSetter({{.FirstLetter}})
}

`
<|endoftext|>