<commit_before>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"io"
	"time"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils/backoff"
	"github.com/golang/protobuf/ptypes/empty"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// GatewayStream interface
type GatewayStream interface {
	Close()
}

type gatewayStream struct {
	closing bool
	ctx     api.Logger
	client  RouterClientForGateway
}

// DefaultBufferSize indicates the default send and receive buffer sizes
var DefaultBufferSize = 10

// GatewayStatusStream for sending gateway statuses
type GatewayStatusStream interface {
	GatewayStream
	Send(*gateway.Status) error
}

// NewMonitoredGatewayStatusStream starts and monitors a GatewayStatusStream
func NewMonitoredGatewayStatusStream(client RouterClientForGateway) GatewayStatusStream {
	s := &gatewayStatusStream{
		ch:  make(chan *gateway.Status, DefaultBufferSize),
		err: make(chan error),
	}
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var retries int

		for {
			// Session channels
			ch := make(chan *gateway.Status)
			errCh := make(chan error)

			// Session client
			client, err := s.client.GatewayStatus()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped GatewayStatus stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start GatewayStatus stream, retrying...")
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Debug("Started GatewayStatus stream")

			// Receive errors
			go func() {
				empty := new(empty.Empty)
				if err := client.RecvMsg(empty); err != nil {
					errCh <- err
				}
				close(errCh)
			}()

			// Send
			go func() {
				for status := range ch {
					s.ctx.Debug("Sending GatewayStatus message")
					if err := client.Send(status); err != nil {
						s.ctx.WithError(err).Warn("Error sending GatewayStatus message")
						break
					}
				}
			}()

			// Monitoring
			var mErr error

		monitor:
			for {
				select {
				case mErr = <-errCh:
					break monitor
				case msg, ok := <-s.ch:
					if !ok {
						break monitor // channel closed
					}
					ch <- msg
				}
			}

			close(ch)
			client.CloseAndRecv()

			if mErr == nil || mErr == io.EOF || grpc.Code(mErr) == codes.Canceled {
				s.ctx.Debug("Stopped GatewayStatus stream")
			} else {
				s.ctx.WithError(mErr).Warn("Error in GatewayStatus stream")
			}

			if s.closing {
				break
			}

			time.Sleep(backoff.Backoff(retries))
			retries++
		}
	}()

	return s
}

type gatewayStatusStream struct {
	gatewayStream
	ch  chan *gateway.Status
	err chan error
}

func (s *gatewayStatusStream) Send(status *gateway.Status) error {
	select {
	case s.ch <- status:
	default:
		s.ctx.Warn("Dropping GatewayStatus message, buffer full")
	}
	return nil
}

func (s *gatewayStatusStream) Close() {
	s.closing = true
	close(s.ch)
}

// UplinkStream for sending uplink messages
type UplinkStream interface {
	GatewayStream
	Send(*UplinkMessage) error
}

// NewMonitoredUplinkStream starts and monitors a UplinkStream
func NewMonitoredUplinkStream(client RouterClientForGateway) UplinkStream {
	s := &uplinkStream{
		ch:  make(chan *UplinkMessage, DefaultBufferSize),
		err: make(chan error),
	}
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var retries int

		for {
			// Session channels
			ch := make(chan *UplinkMessage)
			errCh := make(chan error)

			// Session client
			client, err := s.client.Uplink()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped Uplink stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start Uplink stream, retrying...")
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Debug("Started Uplink stream")

			// Receive errors
			go func() {
				empty := new(empty.Empty)
				if err := client.RecvMsg(empty); err != nil {
					errCh <- err
				}
				close(errCh)
			}()

			// Send
			go func() {
				for message := range ch {
					s.ctx.Debug("Sending Uplink message")
					if err := client.Send(message); err != nil {
						s.ctx.WithError(err).Warn("Error sending Uplink message")
						break
					}
				}
			}()

			// Monitoring
			var mErr error

		monitor:
			for {
				select {
				case mErr = <-errCh:
					break monitor
				case msg, ok := <-s.ch:
					if !ok {
						break monitor // channel closed
					}
					ch <- msg
				}
			}

			close(ch)
			client.CloseAndRecv()

			if mErr == nil || mErr == io.EOF || grpc.Code(mErr) == codes.Canceled {
				s.ctx.Debug("Stopped Uplink stream")
			} else {
				s.ctx.WithError(mErr).Warn("Error in Uplink stream")
			}

			if s.closing {
				break
			}

			time.Sleep(backoff.Backoff(retries))
			retries++
		}
	}()

	return s
}

type uplinkStream struct {
	gatewayStream
	ch  chan *UplinkMessage
	err chan error
}

func (s *uplinkStream) Send(message *UplinkMessage) error {
	select {
	case s.ch <- message:
	default:
		s.ctx.Warn("Dropping Uplink message, buffer full")
	}
	return nil
}

func (s *uplinkStream) Close() {
	s.closing = true
	close(s.ch)
}

// DownlinkStream for sending downlink messages
type DownlinkStream interface {
	GatewayStream
	Channel() <-chan *DownlinkMessage
}

// NewMonitoredDownlinkStream starts and monitors a DownlinkStream
func NewMonitoredDownlinkStream(client RouterClientForGateway) DownlinkStream {
	s := &downlinkStream{
		ch:  make(chan *DownlinkMessage, DefaultBufferSize),
		err: make(chan error),
	}
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var client Router_SubscribeClient
		var err error
		var retries int
		var message *DownlinkMessage

		for {
			client, s.cancel, err = s.client.Subscribe()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped Downlink stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start Downlink stream, retrying...")
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Debug("Started Downlink stream")

			for {
				message, err = client.Recv()
				if message != nil {
					s.ctx.Debug("Receiving Downlink message")
					select {
					case s.ch <- message:
					default:
						s.ctx.Warn("Dropping Downlink message, buffer full")
					}
				}
				if err != nil {
					break
				}
			}

			if err == nil || err == io.EOF || grpc.Code(err) == codes.Canceled {
				s.ctx.Debug("Stopped Downlink stream")
			} else {
				s.ctx.WithError(err).Warn("Error in Downlink stream")
			}

			if s.closing {
				break
			}

			time.Sleep(backoff.Backoff(retries))
			retries++
		}

		close(s.ch)
	}()
	return s
}

type downlinkStream struct {
	gatewayStream
	cancel context.CancelFunc
	ch     chan *DownlinkMessage
	err    chan error
}

func (s *downlinkStream) Close() {
	s.closing = true
	if s.cancel != nil {
		s.cancel()
	}
}

func (s *downlinkStream) Channel() <-chan *DownlinkMessage {
	return s.ch
}
<commit_msg>Wait for stream setup to finish before closing<commit_after>// Copyright © 2016 The Things Network
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

package router

import (
	"io"
	"sync"
	"time"

	"github.com/TheThingsNetwork/ttn/api"
	"github.com/TheThingsNetwork/ttn/api/gateway"
	"github.com/TheThingsNetwork/ttn/utils/backoff"
	"github.com/golang/protobuf/ptypes/empty"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

// GatewayStream interface
type GatewayStream interface {
	Close()
}

type gatewayStream struct {
	closing bool
	setup   sync.WaitGroup
	ctx     api.Logger
	client  RouterClientForGateway
}

// DefaultBufferSize indicates the default send and receive buffer sizes
var DefaultBufferSize = 10

// GatewayStatusStream for sending gateway statuses
type GatewayStatusStream interface {
	GatewayStream
	Send(*gateway.Status) error
}

// NewMonitoredGatewayStatusStream starts and monitors a GatewayStatusStream
func NewMonitoredGatewayStatusStream(client RouterClientForGateway) GatewayStatusStream {
	s := &gatewayStatusStream{
		ch:  make(chan *gateway.Status, DefaultBufferSize),
		err: make(chan error),
	}
	s.setup.Add(1)
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var retries int

		for {
			// Session channels
			ch := make(chan *gateway.Status)
			errCh := make(chan error)

			// Session client
			client, err := s.client.GatewayStatus()
			s.setup.Done()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped GatewayStatus stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start GatewayStatus stream, retrying...")
				s.setup.Add(1)
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Info("Started GatewayStatus stream")

			// Receive errors
			go func() {
				empty := new(empty.Empty)
				if err := client.RecvMsg(empty); err != nil {
					errCh <- err
				}
				close(errCh)
			}()

			// Send
			go func() {
				for status := range ch {
					s.ctx.Debug("Sending GatewayStatus message")
					if err := client.Send(status); err != nil {
						s.ctx.WithError(err).Warn("Error sending GatewayStatus message")
						break
					}
				}
			}()

			// Monitoring
			var mErr error

		monitor:
			for {
				select {
				case mErr = <-errCh:
					break monitor
				case msg, ok := <-s.ch:
					if !ok {
						break monitor // channel closed
					}
					ch <- msg
				}
			}

			close(ch)
			client.CloseAndRecv()

			if mErr == nil || mErr == io.EOF || grpc.Code(mErr) == codes.Canceled {
				s.ctx.Debug("Stopped GatewayStatus stream")
			} else {
				s.ctx.WithError(mErr).Warn("Error in GatewayStatus stream")
			}

			if s.closing {
				break
			}

			s.setup.Add(1)
			time.Sleep(backoff.Backoff(retries))
			retries++
		}
	}()

	return s
}

type gatewayStatusStream struct {
	gatewayStream
	ch  chan *gateway.Status
	err chan error
}

func (s *gatewayStatusStream) Send(status *gateway.Status) error {
	select {
	case s.ch <- status:
	default:
		s.ctx.Warn("Dropping GatewayStatus message, buffer full")
	}
	return nil
}

func (s *gatewayStatusStream) Close() {
	s.setup.Wait()
	s.ctx.Debug("Closing GatewayStatus stream")
	s.closing = true
	close(s.ch)
}

// UplinkStream for sending uplink messages
type UplinkStream interface {
	GatewayStream
	Send(*UplinkMessage) error
}

// NewMonitoredUplinkStream starts and monitors a UplinkStream
func NewMonitoredUplinkStream(client RouterClientForGateway) UplinkStream {
	s := &uplinkStream{
		ch:  make(chan *UplinkMessage, DefaultBufferSize),
		err: make(chan error),
	}
	s.setup.Add(1)
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var retries int

		for {
			// Session channels
			ch := make(chan *UplinkMessage)
			errCh := make(chan error)

			// Session client
			client, err := s.client.Uplink()
			s.setup.Done()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped Uplink stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start Uplink stream, retrying...")
				s.setup.Add(1)
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Info("Started Uplink stream")

			// Receive errors
			go func() {
				empty := new(empty.Empty)
				if err := client.RecvMsg(empty); err != nil {
					errCh <- err
				}
				close(errCh)
			}()

			// Send
			go func() {
				for message := range ch {
					s.ctx.Debug("Sending Uplink message")
					if err := client.Send(message); err != nil {
						s.ctx.WithError(err).Warn("Error sending Uplink message")
						break
					}
				}
			}()

			// Monitoring
			var mErr error

		monitor:
			for {
				select {
				case mErr = <-errCh:
					break monitor
				case msg, ok := <-s.ch:
					if !ok {
						break monitor // channel closed
					}
					ch <- msg
				}
			}

			close(ch)
			client.CloseAndRecv()

			if mErr == nil || mErr == io.EOF || grpc.Code(mErr) == codes.Canceled {
				s.ctx.Debug("Stopped Uplink stream")
			} else {
				s.ctx.WithError(mErr).Warn("Error in Uplink stream")
			}

			if s.closing {
				break
			}

			s.setup.Add(1)
			time.Sleep(backoff.Backoff(retries))
			retries++
		}
	}()

	return s
}

type uplinkStream struct {
	gatewayStream
	ch  chan *UplinkMessage
	err chan error
}

func (s *uplinkStream) Send(message *UplinkMessage) error {
	select {
	case s.ch <- message:
	default:
		s.ctx.Warn("Dropping Uplink message, buffer full")
	}
	return nil
}

func (s *uplinkStream) Close() {
	s.setup.Wait()
	s.ctx.Debug("Closing Uplink stream")
	s.closing = true
	close(s.ch)
}

// DownlinkStream for sending downlink messages
type DownlinkStream interface {
	GatewayStream
	Channel() <-chan *DownlinkMessage
}

// NewMonitoredDownlinkStream starts and monitors a DownlinkStream
func NewMonitoredDownlinkStream(client RouterClientForGateway) DownlinkStream {
	s := &downlinkStream{
		ch:  make(chan *DownlinkMessage, DefaultBufferSize),
		err: make(chan error),
	}
	s.setup.Add(1)
	s.client = client
	s.ctx = client.GetLogger()

	go func() {
		var client Router_SubscribeClient
		var err error
		var retries int
		var message *DownlinkMessage

		for {
			client, s.cancel, err = s.client.Subscribe()
			s.setup.Done()
			if err != nil {
				if grpc.Code(err) == codes.Canceled {
					s.ctx.Debug("Stopped Downlink stream")
					break
				}
				s.ctx.WithError(err).Warn("Could not start Downlink stream, retrying...")
				s.setup.Add(1)
				time.Sleep(backoff.Backoff(retries))
				retries++
				continue
			}
			retries = 0

			s.ctx.Info("Started Downlink stream")

			for {
				message, err = client.Recv()
				if message != nil {
					s.ctx.Debug("Receiving Downlink message")
					select {
					case s.ch <- message:
					default:
						s.ctx.Warn("Dropping Downlink message, buffer full")
					}
				}
				if err != nil {
					break
				}
			}

			if err == nil || err == io.EOF || grpc.Code(err) == codes.Canceled {
				s.ctx.Debug("Stopped Downlink stream")
			} else {
				s.ctx.WithError(err).Warn("Error in Downlink stream")
			}

			if s.closing {
				break
			}

			s.setup.Add(1)
			time.Sleep(backoff.Backoff(retries))
			retries++
		}

		close(s.ch)
	}()
	return s
}

type downlinkStream struct {
	gatewayStream
	cancel context.CancelFunc
	ch     chan *DownlinkMessage
	err    chan error
}

func (s *downlinkStream) Close() {
	s.setup.Wait()
	s.ctx.Debug("Closing Downlink stream")
	s.closing = true
	if s.cancel != nil {
		s.cancel()
	}
}

func (s *downlinkStream) Channel() <-chan *DownlinkMessage {
	return s.ch
}
<|endoftext|>


<commit_before>package routines

import (
	"strconv"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	saramaC "github.com/d33d33/sarama-cluster"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/spf13/viper"

	"github.com/runabove/metronome/src/metronome/kafka"
	"github.com/runabove/metronome/src/metronome/models"
)

// Partition handle a topic partition
type Partition struct {
	Partition int32
	Tasks     chan models.Task
}

// TaskConsumer handle the internal states of the consumer
type TaskConsumer struct {
	client   *saramaC.Client
	consumer *saramaC.Consumer
	drained  bool
	drainWg  sync.WaitGroup
	// group tasks by partition
	partitions     map[int32]chan models.Task
	partitionsChan chan Partition
	hwm            map[int32]int64
	// metrics
	taskCounter              *prometheus.CounterVec
	taskUnprocessableCounter *prometheus.CounterVec
}

// NewTaskComsumer return a new task consumer
func NewTaskComsumer() (*TaskConsumer, error) {
	brokers := viper.GetStringSlice("kafka.brokers")

	config := saramaC.NewConfig()
	config.Config = *kafka.NewConfig()
	config.ClientID = "metronome-scheduler"
	config.Consumer.Offsets.Initial = sarama.OffsetOldest
	config.Group.Return.Notifications = true

	client, err := saramaC.NewClient(brokers, config)
	if err != nil {
		return nil, err
	}

	consumer, err := saramaC.NewConsumerFromClient(client, kafka.GroupSchedulers(), []string{kafka.TopicTasks()})
	if err != nil {
		return nil, err
	}

	tc := &TaskConsumer{
		client:         client,
		consumer:       consumer,
		partitions:     make(map[int32]chan models.Task),
		partitionsChan: make(chan Partition),
	}
	tc.taskCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "metronome",
		Subsystem: "scheduler",
		Name:      "tasks",
		Help:      "Number of tasks processed.",
	},
		[]string{"partition"})
	prometheus.MustRegister(tc.taskCounter)
	tc.taskUnprocessableCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "metronome",
		Subsystem: "scheduler",
		Name:      "tasks_unprocessable",
		Help:      "Number of unprocessable tasks.",
	},
		[]string{"partition"})

	tc.hwm = <-tc.highWaterMarks()
	offsets := make(map[int32]int64)
	messages := 0

	tc.drainWg.Add(1)

	// Progress display
	ticker := time.NewTicker(500 * time.Millisecond)

	go func() {
		for {
			select {
			case msg, ok := <-consumer.Messages():
				if !ok { // shuting down
					return
				}

				// skip if we have already processed this message
				// hapenned at rebalance
				if offsets[msg.Partition] < msg.Offset {
					messages++
					tc.handleMsg(msg)
					offsets[msg.Partition] = msg.Offset
				}

				if !tc.drained && tc.isDrained(tc.hwm, offsets) {
					ticker.Stop()
					tc.drained = true
					tc.drainWg.Done()
				}

			case notif := <-consumer.Notifications():
				log.Infof("Rebalance - claim %v, release %v", notif.Claimed[kafka.TopicTasks()], notif.Released[kafka.TopicTasks()])
				for _, p := range notif.Released[kafka.TopicTasks()] {
					if tc.partitions[p] != nil {
						close(tc.partitions[p])
						delete(tc.partitions, p)
						offsets[p] = 0
					}
				}
				tc.hwm = <-tc.highWaterMarks()
				for _, p := range notif.Claimed[kafka.TopicTasks()] {
					if tc.drained {
						tc.drained = false
						tc.drainWg.Add(1)
					}

					tc.partitions[p] = make(chan models.Task)
					tc.partitionsChan <- Partition{p, tc.partitions[p]}
				}
			case <-ticker.C:
				log.WithField("count", messages).Debug("Loading tasks")
			}
		}
	}()

	return tc, nil
}

// Partitons return the incomming partition channel
func (tc *TaskConsumer) Partitons() <-chan Partition {
	return tc.partitionsChan
}

// WaitForDrain wait for consumer to EOF partitions
func (tc *TaskConsumer) WaitForDrain() {
	tc.drainWg.Wait()
}

// Close the task consumer
func (tc *TaskConsumer) Close() (err error) {
	if e := tc.consumer.Close(); e != nil {
		err = e
	}
	if e := tc.client.Close(); e != nil {
		err = e
	}
	for _, p := range tc.partitions {
		close(p)
	}
	if !tc.drained {
		tc.drainWg.Done()
	}
	return
}

// Handle incomming messages
func (tc *TaskConsumer) handleMsg(msg *sarama.ConsumerMessage) {
	tc.taskCounter.WithLabelValues(strconv.Itoa(int(msg.Partition))).Inc()
	var t models.Task
	if err := t.FromKafka(msg); err != nil {
		tc.taskUnprocessableCounter.WithLabelValues(strconv.Itoa(int(msg.Partition))).Inc()
		log.Error(err)
		return
	}
	log.Debugf("Task received: %v partition %v", t.ToJSON(), msg.Partition)
	tc.partitions[msg.Partition] <- t
}

// Retrieve highWaterMarks for each partition
func (tc *TaskConsumer) highWaterMarks() chan map[int32]int64 {
	resChan := make(chan map[int32]int64)

	go func() {
		for {
			parts, err := tc.client.Partitions(kafka.TopicTasks())
			if err != nil {
				log.Warn("Can't get topic. Retry")
				continue
			}

			res := make(map[int32]int64)
			for p := range parts {
				i, err := tc.client.GetOffset(kafka.TopicTasks(), int32(p), sarama.OffsetNewest)
				if err != nil {
					log.Panic(err)
				}

				res[int32(p)] = i
			}

			resChan <- res
			close(resChan)
			break
		}
	}()

	return resChan
}

// Check if consumer reach EOF on all the partitions
func (tc *TaskConsumer) isDrained(hwm, offsets map[int32]int64) bool {
	subs := tc.consumer.Subscriptions()[kafka.TopicTasks()]

	for partition := range subs {
		part := int32(partition)
		if _, ok := hwm[part]; !ok {
			log.Panicf("Missing HighWaterMarks for partition %v", part)
		}
		if hwm[part] == 0 {
			continue
		}
		// No message received for partiton
		if _, ok := offsets[part]; !ok {
			return false
		}
		// Check offset
		if (offsets[part] + 1) < hwm[part] {
			return false
		}
	}

	return true
}
<commit_msg>fix(scheduler): first message is not consumed<commit_after>package routines

import (
	"strconv"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	log "github.com/Sirupsen/logrus"
	saramaC "github.com/d33d33/sarama-cluster"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/spf13/viper"

	"github.com/runabove/metronome/src/metronome/kafka"
	"github.com/runabove/metronome/src/metronome/models"
)

// Partition handle a topic partition
type Partition struct {
	Partition int32
	Tasks     chan models.Task
}

// TaskConsumer handle the internal states of the consumer
type TaskConsumer struct {
	client   *saramaC.Client
	consumer *saramaC.Consumer
	drained  bool
	drainWg  sync.WaitGroup
	// group tasks by partition
	partitions     map[int32]chan models.Task
	partitionsChan chan Partition
	hwm            map[int32]int64
	// metrics
	taskCounter              *prometheus.CounterVec
	taskUnprocessableCounter *prometheus.CounterVec
}

// NewTaskComsumer return a new task consumer
func NewTaskComsumer() (*TaskConsumer, error) {
	brokers := viper.GetStringSlice("kafka.brokers")

	config := saramaC.NewConfig()
	config.Config = *kafka.NewConfig()
	config.ClientID = "metronome-scheduler"
	config.Consumer.Offsets.Initial = sarama.OffsetOldest
	config.Group.Return.Notifications = true

	client, err := saramaC.NewClient(brokers, config)
	if err != nil {
		return nil, err
	}

	consumer, err := saramaC.NewConsumerFromClient(client, kafka.GroupSchedulers(), []string{kafka.TopicTasks()})
	if err != nil {
		return nil, err
	}

	tc := &TaskConsumer{
		client:         client,
		consumer:       consumer,
		partitions:     make(map[int32]chan models.Task),
		partitionsChan: make(chan Partition),
	}
	tc.taskCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "metronome",
		Subsystem: "scheduler",
		Name:      "tasks",
		Help:      "Number of tasks processed.",
	},
		[]string{"partition"})
	prometheus.MustRegister(tc.taskCounter)
	tc.taskUnprocessableCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
		Namespace: "metronome",
		Subsystem: "scheduler",
		Name:      "tasks_unprocessable",
		Help:      "Number of unprocessable tasks.",
	},
		[]string{"partition"})

	tc.hwm = <-tc.highWaterMarks()
	offsets := make(map[int32]int64)
	messages := 0

	// init offsets
	for p := range tc.hwm {
		offsets[p] = -1
	}

	tc.drainWg.Add(1)

	// Progress display
	ticker := time.NewTicker(500 * time.Millisecond)

	go func() {
		for {
			select {
			case msg, ok := <-consumer.Messages():
				if !ok { // shuting down
					return
				}

				// skip if we have already processed this message
				// hapenned at rebalance
				if offsets[msg.Partition] < msg.Offset {
					messages++
					tc.handleMsg(msg)
					offsets[msg.Partition] = msg.Offset
				}

				if !tc.drained && tc.isDrained(tc.hwm, offsets) {
					ticker.Stop()
					tc.drained = true
					tc.drainWg.Done()
				}

			case notif := <-consumer.Notifications():
				log.Infof("Rebalance - claim %v, release %v", notif.Claimed[kafka.TopicTasks()], notif.Released[kafka.TopicTasks()])
				for _, p := range notif.Released[kafka.TopicTasks()] {
					if tc.partitions[p] != nil {
						close(tc.partitions[p])
						delete(tc.partitions, p)
						offsets[p] = 0
					}
				}
				tc.hwm = <-tc.highWaterMarks()
				for _, p := range notif.Claimed[kafka.TopicTasks()] {
					if tc.drained {
						tc.drained = false
						tc.drainWg.Add(1)
					}

					tc.partitions[p] = make(chan models.Task)
					tc.partitionsChan <- Partition{p, tc.partitions[p]}
				}
			case <-ticker.C:
				log.WithField("count", messages).Debug("Loading tasks")
			}
		}
	}()

	return tc, nil
}

// Partitons return the incomming partition channel
func (tc *TaskConsumer) Partitons() <-chan Partition {
	return tc.partitionsChan
}

// WaitForDrain wait for consumer to EOF partitions
func (tc *TaskConsumer) WaitForDrain() {
	tc.drainWg.Wait()
}

// Close the task consumer
func (tc *TaskConsumer) Close() (err error) {
	if e := tc.consumer.Close(); e != nil {
		err = e
	}
	if e := tc.client.Close(); e != nil {
		err = e
	}
	for _, p := range tc.partitions {
		close(p)
	}
	if !tc.drained {
		tc.drainWg.Done()
	}
	return
}

// Handle incomming messages
func (tc *TaskConsumer) handleMsg(msg *sarama.ConsumerMessage) {
	tc.taskCounter.WithLabelValues(strconv.Itoa(int(msg.Partition))).Inc()
	var t models.Task
	if err := t.FromKafka(msg); err != nil {
		tc.taskUnprocessableCounter.WithLabelValues(strconv.Itoa(int(msg.Partition))).Inc()
		log.Error(err)
		return
	}
	log.Debugf("Task received: %v partition %v", t.ToJSON(), msg.Partition)
	tc.partitions[msg.Partition] <- t
}

// Retrieve highWaterMarks for each partition
func (tc *TaskConsumer) highWaterMarks() chan map[int32]int64 {
	resChan := make(chan map[int32]int64)

	go func() {
		for {
			parts, err := tc.client.Partitions(kafka.TopicTasks())
			if err != nil {
				log.Warn("Can't get topic. Retry")
				continue
			}

			res := make(map[int32]int64)
			for p := range parts {
				i, err := tc.client.GetOffset(kafka.TopicTasks(), int32(p), sarama.OffsetNewest)
				if err != nil {
					log.Panic(err)
				}

				res[int32(p)] = i
			}

			resChan <- res
			close(resChan)
			break
		}
	}()

	return resChan
}

// Check if consumer reach EOF on all the partitions
func (tc *TaskConsumer) isDrained(hwm, offsets map[int32]int64) bool {
	subs := tc.consumer.Subscriptions()[kafka.TopicTasks()]

	for partition := range subs {
		part := int32(partition)
		if _, ok := hwm[part]; !ok {
			log.Panicf("Missing HighWaterMarks for partition %v", part)
		}
		if hwm[part] == 0 {
			continue
		}
		// No message received for partiton
		if _, ok := offsets[part]; !ok {
			return false
		}
		// Check offset
		if (offsets[part] + 1) < hwm[part] {
			return false
		}
	}

	return true
}
<|endoftext|>


<commit_before>// Copyright (c) 2019 Minoru Osuka
//
// 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 grpc

import (
	"context"
	"errors"
	"math"
	"time"

	"github.com/blevesearch/bleve"
	"github.com/golang/protobuf/ptypes/any"
	"github.com/golang/protobuf/ptypes/empty"
	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
	blasterrors "github.com/mosuka/blast/errors"
	"github.com/mosuka/blast/protobuf"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type Client struct {
	ctx    context.Context
	cancel context.CancelFunc
	conn   *grpc.ClientConn
	client protobuf.BlastClient
}

func NewContext() (context.Context, context.CancelFunc) {
	baseCtx := context.TODO()
	return context.WithTimeout(baseCtx, 60*time.Second)
}

func NewClient(address string) (*Client, error) {
	ctx, cancel := NewContext()

	streamRetryOpts := []grpc_retry.CallOption{
		grpc_retry.Disable(),
	}

	unaryRetryOpts := []grpc_retry.CallOption{
		grpc_retry.WithBackoff(grpc_retry.BackoffLinear(100 * time.Millisecond)),
		grpc_retry.WithCodes(codes.Unavailable),
		grpc_retry.WithMax(100),
	}

	dialOpts := []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithDefaultCallOptions(
			grpc.MaxCallSendMsgSize(math.MaxInt32),
			grpc.MaxCallRecvMsgSize(math.MaxInt32),
		),
		grpc.WithStreamInterceptor(grpc_retry.StreamClientInterceptor(streamRetryOpts...)),
		grpc.WithUnaryInterceptor(grpc_retry.UnaryClientInterceptor(unaryRetryOpts...)),
	}

	conn, err := grpc.DialContext(ctx, address, dialOpts...)
	if err != nil {
		return nil, err
	}

	return &Client{
		ctx:    ctx,
		cancel: cancel,
		conn:   conn,
		client: protobuf.NewBlastClient(conn),
	}, nil
}

func (c *Client) Cancel() {
	c.cancel()
}

func (c *Client) Close() error {
	c.Cancel()
	if c.conn != nil {
		return c.conn.Close()
	}

	return c.ctx.Err()
}

func (c *Client) GetAddress() string {
	return c.conn.Target()
}

func (c *Client) GetNode(id string, opts ...grpc.CallOption) (map[string]interface{}, error) {
	req := &protobuf.GetNodeRequest{
		Id: id,
	}

	resp, err := c.client.GetNode(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	ins, err := protobuf.MarshalAny(resp.Metadata)
	metadata := *ins.(*map[string]interface{})

	node := map[string]interface{}{
		"metadata": metadata,
		"state":    resp.State,
	}

	return node, nil
}

func (c *Client) SetNode(id string, metadata map[string]interface{}, opts ...grpc.CallOption) error {
	metadataAny := &any.Any{}
	err := protobuf.UnmarshalAny(metadata, metadataAny)
	if err != nil {
		return err
	}

	req := &protobuf.SetNodeRequest{
		Id:       id,
		Metadata: metadataAny,
	}

	_, err = c.client.SetNode(c.ctx, req, opts...)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) DeleteNode(id string, opts ...grpc.CallOption) error {
	req := &protobuf.DeleteNodeRequest{
		Id: id,
	}

	_, err := c.client.DeleteNode(c.ctx, req, opts...)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) GetCluster(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetCluster(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	ins, err := protobuf.MarshalAny(resp.Cluster)
	cluster := *ins.(*map[string]interface{})

	return cluster, nil
}

func (c *Client) WatchCluster(opts ...grpc.CallOption) (protobuf.Blast_WatchClusterClient, error) {
	req := &empty.Empty{}

	watchClient, err := c.client.WatchCluster(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)
		return nil, errors.New(st.Message())
	}

	return watchClient, nil
}

func (c *Client) Snapshot(opts ...grpc.CallOption) error {
	_, err := c.client.Snapshot(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return errors.New(st.Message())
	}

	return nil
}

func (c *Client) LivenessProbe(opts ...grpc.CallOption) (string, error) {
	resp, err := c.client.LivenessProbe(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return protobuf.LivenessProbeResponse_UNKNOWN.String(), errors.New(st.Message())
	}

	return resp.State.String(), nil
}

func (c *Client) ReadinessProbe(opts ...grpc.CallOption) (string, error) {
	resp, err := c.client.ReadinessProbe(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return protobuf.ReadinessProbeResponse_UNKNOWN.String(), errors.New(st.Message())
	}

	return resp.State.String(), nil
}

func (c *Client) GetState(key string, opts ...grpc.CallOption) (interface{}, error) {
	req := &protobuf.GetStateRequest{
		Key: key,
	}

	resp, err := c.client.GetState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return nil, blasterrors.ErrNotFound
		default:
			return nil, errors.New(st.Message())
		}
	}

	value, err := protobuf.MarshalAny(resp.Value)

	return value, nil
}

func (c *Client) SetState(key string, value interface{}, opts ...grpc.CallOption) error {
	valueAny := &any.Any{}
	err := protobuf.UnmarshalAny(value, valueAny)
	if err != nil {
		return err
	}

	req := &protobuf.SetStateRequest{
		Key:   key,
		Value: valueAny,
	}

	_, err = c.client.SetState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return blasterrors.ErrNotFound
		default:
			return errors.New(st.Message())
		}
	}

	return nil
}

func (c *Client) DeleteState(key string, opts ...grpc.CallOption) error {
	req := &protobuf.DeleteStateRequest{
		Key: key,
	}

	_, err := c.client.DeleteState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return blasterrors.ErrNotFound
		default:
			return errors.New(st.Message())
		}
	}

	return nil
}

func (c *Client) WatchState(key string, opts ...grpc.CallOption) (protobuf.Blast_WatchStateClient, error) {
	req := &protobuf.WatchStateRequest{
		Key: key,
	}

	watchClient, err := c.client.WatchState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)
		return nil, errors.New(st.Message())
	}

	return watchClient, nil
}

func (c *Client) GetDocument(id string, opts ...grpc.CallOption) (map[string]interface{}, error) {
	req := &protobuf.GetDocumentRequest{
		Id: id,
	}

	resp, err := c.client.GetDocument(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return nil, blasterrors.ErrNotFound
		default:
			return nil, errors.New(st.Message())
		}
	}

	ins, err := protobuf.MarshalAny(resp.Fields)
	fields := *ins.(*map[string]interface{})

	return fields, nil
}

func (c *Client) Search(searchRequest *bleve.SearchRequest, opts ...grpc.CallOption) (*bleve.SearchResult, error) {
	// bleve.SearchRequest -> Any
	searchRequestAny := &any.Any{}
	err := protobuf.UnmarshalAny(searchRequest, searchRequestAny)
	if err != nil {
		return nil, err
	}

	req := &protobuf.SearchRequest{
		SearchRequest: searchRequestAny,
	}

	resp, err := c.client.Search(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	// Any -> bleve.SearchResult
	searchResultInstance, err := protobuf.MarshalAny(resp.SearchResult)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	if searchResultInstance == nil {
		return nil, errors.New("nil")
	}
	searchResult := searchResultInstance.(*bleve.SearchResult)

	return searchResult, nil
}

func (c *Client) IndexDocument(docs []map[string]interface{}, opts ...grpc.CallOption) (int, error) {
	stream, err := c.client.IndexDocument(c.ctx, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return -1, errors.New(st.Message())
	}

	for _, doc := range docs {
		id := doc["id"].(string)
		fields := doc["fields"].(map[string]interface{})

		fieldsAny := &any.Any{}
		err := protobuf.UnmarshalAny(&fields, fieldsAny)
		if err != nil {
			return -1, err
		}

		req := &protobuf.IndexDocumentRequest{
			Id:     id,
			Fields: fieldsAny,
		}

		err = stream.Send(req)
		if err != nil {
			return -1, err
		}
	}

	resp, err := stream.CloseAndRecv()
	if err != nil {
		return -1, err
	}

	return int(resp.Count), nil
}

func (c *Client) DeleteDocument(ids []string, opts ...grpc.CallOption) (int, error) {
	stream, err := c.client.DeleteDocument(c.ctx, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return -1, errors.New(st.Message())
	}

	for _, id := range ids {
		req := &protobuf.DeleteDocumentRequest{
			Id: id,
		}

		err := stream.Send(req)
		if err != nil {
			return -1, err
		}
	}

	resp, err := stream.CloseAndRecv()
	if err != nil {
		return -1, err
	}

	return int(resp.Count), nil
}

func (c *Client) GetIndexConfig(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetIndexConfig(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	indexConfigIntr, err := protobuf.MarshalAny(resp.IndexConfig)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	indexConfig := *indexConfigIntr.(*map[string]interface{})

	return indexConfig, nil
}

func (c *Client) GetIndexStats(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetIndexStats(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	indexStatsIntr, err := protobuf.MarshalAny(resp.IndexStats)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	indexStats := *indexStatsIntr.(*map[string]interface{})

	return indexStats, nil
}
<commit_msg>Fix bug<commit_after>// Copyright (c) 2019 Minoru Osuka
//
// 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 grpc

import (
	"context"
	"errors"
	"math"
	"time"

	"github.com/blevesearch/bleve"
	"github.com/golang/protobuf/ptypes/any"
	"github.com/golang/protobuf/ptypes/empty"
	grpc_retry "github.com/grpc-ecosystem/go-grpc-middleware/retry"
	blasterrors "github.com/mosuka/blast/errors"
	"github.com/mosuka/blast/protobuf"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

type Client struct {
	ctx    context.Context
	cancel context.CancelFunc
	conn   *grpc.ClientConn
	client protobuf.BlastClient
}

func NewContext() (context.Context, context.CancelFunc) {
	baseCtx := context.TODO()
	//return context.WithTimeout(baseCtx, 60*time.Second)
	return context.WithCancel(baseCtx)
}

func NewClient(address string) (*Client, error) {
	ctx, cancel := NewContext()

	streamRetryOpts := []grpc_retry.CallOption{
		grpc_retry.Disable(),
	}

	unaryRetryOpts := []grpc_retry.CallOption{
		grpc_retry.WithBackoff(grpc_retry.BackoffLinear(100 * time.Millisecond)),
		grpc_retry.WithCodes(codes.Unavailable),
		grpc_retry.WithMax(100),
	}

	dialOpts := []grpc.DialOption{
		grpc.WithInsecure(),
		grpc.WithDefaultCallOptions(
			grpc.MaxCallSendMsgSize(math.MaxInt32),
			grpc.MaxCallRecvMsgSize(math.MaxInt32),
		),
		grpc.WithStreamInterceptor(grpc_retry.StreamClientInterceptor(streamRetryOpts...)),
		grpc.WithUnaryInterceptor(grpc_retry.UnaryClientInterceptor(unaryRetryOpts...)),
	}

	conn, err := grpc.DialContext(ctx, address, dialOpts...)
	if err != nil {
		return nil, err
	}

	return &Client{
		ctx:    ctx,
		cancel: cancel,
		conn:   conn,
		client: protobuf.NewBlastClient(conn),
	}, nil
}

func (c *Client) Cancel() {
	c.cancel()
}

func (c *Client) Close() error {
	c.Cancel()
	if c.conn != nil {
		return c.conn.Close()
	}

	return c.ctx.Err()
}

func (c *Client) GetAddress() string {
	return c.conn.Target()
}

func (c *Client) GetNode(id string, opts ...grpc.CallOption) (map[string]interface{}, error) {
	req := &protobuf.GetNodeRequest{
		Id: id,
	}

	resp, err := c.client.GetNode(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	ins, err := protobuf.MarshalAny(resp.Metadata)
	metadata := *ins.(*map[string]interface{})

	node := map[string]interface{}{
		"metadata": metadata,
		"state":    resp.State,
	}

	return node, nil
}

func (c *Client) SetNode(id string, metadata map[string]interface{}, opts ...grpc.CallOption) error {
	metadataAny := &any.Any{}
	err := protobuf.UnmarshalAny(metadata, metadataAny)
	if err != nil {
		return err
	}

	req := &protobuf.SetNodeRequest{
		Id:       id,
		Metadata: metadataAny,
	}

	_, err = c.client.SetNode(c.ctx, req, opts...)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) DeleteNode(id string, opts ...grpc.CallOption) error {
	req := &protobuf.DeleteNodeRequest{
		Id: id,
	}

	_, err := c.client.DeleteNode(c.ctx, req, opts...)
	if err != nil {
		return err
	}

	return nil
}

func (c *Client) GetCluster(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetCluster(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	ins, err := protobuf.MarshalAny(resp.Cluster)
	cluster := *ins.(*map[string]interface{})

	return cluster, nil
}

func (c *Client) WatchCluster(opts ...grpc.CallOption) (protobuf.Blast_WatchClusterClient, error) {
	req := &empty.Empty{}

	watchClient, err := c.client.WatchCluster(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)
		return nil, errors.New(st.Message())
	}

	return watchClient, nil
}

func (c *Client) Snapshot(opts ...grpc.CallOption) error {
	_, err := c.client.Snapshot(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return errors.New(st.Message())
	}

	return nil
}

func (c *Client) LivenessProbe(opts ...grpc.CallOption) (string, error) {
	resp, err := c.client.LivenessProbe(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return protobuf.LivenessProbeResponse_UNKNOWN.String(), errors.New(st.Message())
	}

	return resp.State.String(), nil
}

func (c *Client) ReadinessProbe(opts ...grpc.CallOption) (string, error) {
	resp, err := c.client.ReadinessProbe(c.ctx, &empty.Empty{})
	if err != nil {
		st, _ := status.FromError(err)

		return protobuf.ReadinessProbeResponse_UNKNOWN.String(), errors.New(st.Message())
	}

	return resp.State.String(), nil
}

func (c *Client) GetState(key string, opts ...grpc.CallOption) (interface{}, error) {
	req := &protobuf.GetStateRequest{
		Key: key,
	}

	resp, err := c.client.GetState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return nil, blasterrors.ErrNotFound
		default:
			return nil, errors.New(st.Message())
		}
	}

	value, err := protobuf.MarshalAny(resp.Value)

	return value, nil
}

func (c *Client) SetState(key string, value interface{}, opts ...grpc.CallOption) error {
	valueAny := &any.Any{}
	err := protobuf.UnmarshalAny(value, valueAny)
	if err != nil {
		return err
	}

	req := &protobuf.SetStateRequest{
		Key:   key,
		Value: valueAny,
	}

	_, err = c.client.SetState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return blasterrors.ErrNotFound
		default:
			return errors.New(st.Message())
		}
	}

	return nil
}

func (c *Client) DeleteState(key string, opts ...grpc.CallOption) error {
	req := &protobuf.DeleteStateRequest{
		Key: key,
	}

	_, err := c.client.DeleteState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return blasterrors.ErrNotFound
		default:
			return errors.New(st.Message())
		}
	}

	return nil
}

func (c *Client) WatchState(key string, opts ...grpc.CallOption) (protobuf.Blast_WatchStateClient, error) {
	req := &protobuf.WatchStateRequest{
		Key: key,
	}

	watchClient, err := c.client.WatchState(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)
		return nil, errors.New(st.Message())
	}

	return watchClient, nil
}

func (c *Client) GetDocument(id string, opts ...grpc.CallOption) (map[string]interface{}, error) {
	req := &protobuf.GetDocumentRequest{
		Id: id,
	}

	resp, err := c.client.GetDocument(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		switch st.Code() {
		case codes.NotFound:
			return nil, blasterrors.ErrNotFound
		default:
			return nil, errors.New(st.Message())
		}
	}

	ins, err := protobuf.MarshalAny(resp.Fields)
	fields := *ins.(*map[string]interface{})

	return fields, nil
}

func (c *Client) Search(searchRequest *bleve.SearchRequest, opts ...grpc.CallOption) (*bleve.SearchResult, error) {
	// bleve.SearchRequest -> Any
	searchRequestAny := &any.Any{}
	err := protobuf.UnmarshalAny(searchRequest, searchRequestAny)
	if err != nil {
		return nil, err
	}

	req := &protobuf.SearchRequest{
		SearchRequest: searchRequestAny,
	}

	resp, err := c.client.Search(c.ctx, req, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	// Any -> bleve.SearchResult
	searchResultInstance, err := protobuf.MarshalAny(resp.SearchResult)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	if searchResultInstance == nil {
		return nil, errors.New("nil")
	}
	searchResult := searchResultInstance.(*bleve.SearchResult)

	return searchResult, nil
}

func (c *Client) IndexDocument(docs []map[string]interface{}, opts ...grpc.CallOption) (int, error) {
	stream, err := c.client.IndexDocument(c.ctx, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return -1, errors.New(st.Message())
	}

	for _, doc := range docs {
		id := doc["id"].(string)
		fields := doc["fields"].(map[string]interface{})

		fieldsAny := &any.Any{}
		err := protobuf.UnmarshalAny(&fields, fieldsAny)
		if err != nil {
			return -1, err
		}

		req := &protobuf.IndexDocumentRequest{
			Id:     id,
			Fields: fieldsAny,
		}

		err = stream.Send(req)
		if err != nil {
			return -1, err
		}
	}

	resp, err := stream.CloseAndRecv()
	if err != nil {
		return -1, err
	}

	return int(resp.Count), nil
}

func (c *Client) DeleteDocument(ids []string, opts ...grpc.CallOption) (int, error) {
	stream, err := c.client.DeleteDocument(c.ctx, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return -1, errors.New(st.Message())
	}

	for _, id := range ids {
		req := &protobuf.DeleteDocumentRequest{
			Id: id,
		}

		err := stream.Send(req)
		if err != nil {
			return -1, err
		}
	}

	resp, err := stream.CloseAndRecv()
	if err != nil {
		return -1, err
	}

	return int(resp.Count), nil
}

func (c *Client) GetIndexConfig(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetIndexConfig(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	indexConfigIntr, err := protobuf.MarshalAny(resp.IndexConfig)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	indexConfig := *indexConfigIntr.(*map[string]interface{})

	return indexConfig, nil
}

func (c *Client) GetIndexStats(opts ...grpc.CallOption) (map[string]interface{}, error) {
	resp, err := c.client.GetIndexStats(c.ctx, &empty.Empty{}, opts...)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}

	indexStatsIntr, err := protobuf.MarshalAny(resp.IndexStats)
	if err != nil {
		st, _ := status.FromError(err)

		return nil, errors.New(st.Message())
	}
	indexStats := *indexStatsIntr.(*map[string]interface{})

	return indexStats, nil
}
<|endoftext|>


<commit_before>package progress

import (
	"fmt"
	"time"
)

// SimpleProgressE ...
func SimpleProgressE(printChar string, tickInterval time.Duration, action func() error) error {
	var actionError error
	SimpleProgress(printChar, tickInterval, func() {
		actionError = action()
	})
	return actionError
}

// SimpleProgress ...
// action : have to be a synchronous action!
// tickInterval : e.g. : 5000 * time.Millisecond
func SimpleProgress(printChar string, tickInterval time.Duration, action func()) {
	// run async
	finishedChan := make(chan bool)

	go func() {
		action()
		finishedChan <- true
	}()

	fmt.Print(printChar)
	isRunFinished := false
	for !isRunFinished {
		select {
		case <-finishedChan:
			isRunFinished = true
		case <-time.Tick(tickInterval):
			fmt.Print(printChar)
		}
	}
	fmt.Println()
}
<commit_msg>progress - don't print right away<commit_after>package progress

import (
	"fmt"
	"time"
)

// SimpleProgressE ...
func SimpleProgressE(printChar string, tickInterval time.Duration, action func() error) error {
	var actionError error
	SimpleProgress(printChar, tickInterval, func() {
		actionError = action()
	})
	return actionError
}

// SimpleProgress ...
// action : have to be a synchronous action!
// tickInterval : e.g. : 5000 * time.Millisecond
func SimpleProgress(printChar string, tickInterval time.Duration, action func()) {
	// run async
	finishedChan := make(chan bool)

	go func() {
		action()
		finishedChan <- true
	}()

	isRunFinished := false
	for !isRunFinished {
		select {
		case <-finishedChan:
			isRunFinished = true
		case <-time.Tick(tickInterval):
			fmt.Print(printChar)
		}
	}
	fmt.Println()
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"fmt"

	"github.com/WeisswurstSystems/WWM-BB/user/usecase/activate"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/changePassword"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/deleteAccount"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/register"
)

func ExampleCommandHandler_Activate() {
	// Send a json Request in this form
	var request activate.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "registrationID": ""
	// }
}

func ExampleCommandHandler_ChangePassword() {
	// Send a json Request in this form
	var request changePassword.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "login": {
	//     "mail": "",
	//     "password": ""
	//   },
	//   "password": ""
	// }
}

func ExampleCommandHandler_DeleteAccount() {
	// Send a json Request in this form
	var request deleteAccount.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "login": {
	//     "mail": "",
	//     "password": ""
	//   }
	// }
}

func ExampleCommandHandler_Register() {
	// Send a json Request in this form
	var request register.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "mail": "",
	//   "password": "",
	//   "mailEnabled": false
	// }
}
<commit_msg>Update example_test.go<commit_after>package command

import (
	"encoding/json"
	"fmt"

	"github.com/WeisswurstSystems/WWM-BB/user/usecase/activate"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/changePassword"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/deleteAccount"
	"github.com/WeisswurstSystems/WWM-BB/user/usecase/register"
)

func ExampleCommandHandler_Activate() {
	// Send a json Request in this form
	var request activate.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "registrationID": ""
	// }
}

func ExampleCommandHandler_ChangePassword() {
	// Send a json Request in this form
	var request changePassword.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "login": {
	//     "mail": "",
	//     "password": ""
	//   },
	//   "password": "",
	//   "passwordRepeated": ""
	// }
}

func ExampleCommandHandler_DeleteAccount() {
	// Send a json Request in this form
	var request deleteAccount.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "login": {
	//     "mail": "",
	//     "password": ""
	//   }
	// }
}

func ExampleCommandHandler_Register() {
	// Send a json Request in this form
	var request register.Request
	data, _ := json.MarshalIndent(request, "", "  ")
	fmt.Printf("%s", data)

	// Output:
	// {
	//   "mail": "",
	//   "password": "",
	//   "mailEnabled": false
	// }
}
<|endoftext|>


<commit_before>package merkletree

import (
	"bytes"
	"crypto/sha256"
	"testing"
)

// TestReaderRoot calls ReaderRoot on a manually crafted dataset
// and checks the output.
func TestReaderRoot(t *testing.T) {
	mt := CreateMerkleTester(t)
	bytes8 := []byte{0, 1, 2, 3, 4, 5, 6, 7}
	reader := bytes.NewReader(bytes8)
	root, err := ReaderRoot(reader, sha256.New(), 1)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(root, mt.roots[8]) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
}

// TestReaderRootPadding passes ReaderRoot a reader that has too few bytes to
// fill the last segment. The segment should not be padded out.
func TestReaderRootPadding(t *testing.T) {
	bytes1 := []byte{1}
	reader := bytes.NewReader(bytes1)
	root, err := ReaderRoot(reader, sha256.New(), 2)
	if err != nil {
		t.Fatal(err)
	}

	expectedRoot := sum(sha256.New(), []byte{0, 1})
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}

	bytes3 := []byte{1, 2, 3}
	reader = bytes.NewReader(bytes3)
	root, err = ReaderRoot(reader, sha256.New(), 2)
	if err != nil {
		t.Fatal(err)
	}

	baseLeft := sum(sha256.New(), []byte{0, 1, 2})
	baseRight := sum(sha256.New(), []byte{0, 3})
	expectedRoot = sum(sha256.New(), append(append([]byte{1}, baseLeft...), baseRight...))
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
}

// TestBuildReaderProof calls BuildReaderProof on a manually crafted dataset
// and checks the output.
func TestBuilReaderProof(t *testing.T) {
	mt := CreateMerkleTester(t)
	bytes7 := []byte{0, 1, 2, 3, 4, 5, 6}
	reader := bytes.NewReader(bytes7)
	root, proofSet, numLeaves, err := BuildReaderProof(reader, sha256.New(), 1, 5)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(root, mt.roots[7]) != 0 {
		t.Error("BuildReaderProof returned the wrong root")
	}
	if len(proofSet) != len(mt.proofSets[7][5]) {
		t.Fatal("BuildReaderProof returned a proof with the wrong length")
	}
	for i := range proofSet {
		if bytes.Compare(proofSet[i], mt.proofSets[7][5][i]) != 0 {
			t.Error("BuildReaderProof returned an incorrect proof")
		}
	}
	if numLeaves != 7 {
		t.Error("BuildReaderProof returned the wrong number of leaves")
	}
}

// TestBuildReaderProofPadding passes BuildReaderProof a reader that has too
// few bytes to fill the last segment. The segment should not be padded out.
func TestBuildReaderProofPadding(t *testing.T) {
	bytes1 := []byte{1}
	reader := bytes.NewReader(bytes1)
	root, proofSet, numLeaves, err := BuildReaderProof(reader, sha256.New(), 2, 0)
	if err != nil {
		t.Fatal(err)
	}

	expectedRoot := sum(sha256.New(), []byte{0, 1})
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
	if len(proofSet) != 1 {
		t.Fatal("proofSet is the incorrect lenght")
	}
	if bytes.Compare(proofSet[0], []byte{1}) != 0 {
		t.Error("proofSet is incorrect")
	}
	if numLeaves != 1 {
		t.Error("wrong number of leaves returned")
	}
}

// TestEmptyReader passes an empty reader into BuildReaderProof.
func TestEmptyReader(t *testing.T) {
	_, _, _, err := BuildReaderProof(new(bytes.Reader), sha256.New(), 64, 5)
	if err == nil {
		t.Error(err)
	}
}
<commit_msg>correct spelling mistake<commit_after>package merkletree

import (
	"bytes"
	"crypto/sha256"
	"testing"
)

// TestReaderRoot calls ReaderRoot on a manually crafted dataset
// and checks the output.
func TestReaderRoot(t *testing.T) {
	mt := CreateMerkleTester(t)
	bytes8 := []byte{0, 1, 2, 3, 4, 5, 6, 7}
	reader := bytes.NewReader(bytes8)
	root, err := ReaderRoot(reader, sha256.New(), 1)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(root, mt.roots[8]) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
}

// TestReaderRootPadding passes ReaderRoot a reader that has too few bytes to
// fill the last segment. The segment should not be padded out.
func TestReaderRootPadding(t *testing.T) {
	bytes1 := []byte{1}
	reader := bytes.NewReader(bytes1)
	root, err := ReaderRoot(reader, sha256.New(), 2)
	if err != nil {
		t.Fatal(err)
	}

	expectedRoot := sum(sha256.New(), []byte{0, 1})
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}

	bytes3 := []byte{1, 2, 3}
	reader = bytes.NewReader(bytes3)
	root, err = ReaderRoot(reader, sha256.New(), 2)
	if err != nil {
		t.Fatal(err)
	}

	baseLeft := sum(sha256.New(), []byte{0, 1, 2})
	baseRight := sum(sha256.New(), []byte{0, 3})
	expectedRoot = sum(sha256.New(), append(append([]byte{1}, baseLeft...), baseRight...))
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
}

// TestBuildReaderProof calls BuildReaderProof on a manually crafted dataset
// and checks the output.
func TestBuilReaderProof(t *testing.T) {
	mt := CreateMerkleTester(t)
	bytes7 := []byte{0, 1, 2, 3, 4, 5, 6}
	reader := bytes.NewReader(bytes7)
	root, proofSet, numLeaves, err := BuildReaderProof(reader, sha256.New(), 1, 5)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Compare(root, mt.roots[7]) != 0 {
		t.Error("BuildReaderProof returned the wrong root")
	}
	if len(proofSet) != len(mt.proofSets[7][5]) {
		t.Fatal("BuildReaderProof returned a proof with the wrong length")
	}
	for i := range proofSet {
		if bytes.Compare(proofSet[i], mt.proofSets[7][5][i]) != 0 {
			t.Error("BuildReaderProof returned an incorrect proof")
		}
	}
	if numLeaves != 7 {
		t.Error("BuildReaderProof returned the wrong number of leaves")
	}
}

// TestBuildReaderProofPadding passes BuildReaderProof a reader that has too
// few bytes to fill the last segment. The segment should not be padded out.
func TestBuildReaderProofPadding(t *testing.T) {
	bytes1 := []byte{1}
	reader := bytes.NewReader(bytes1)
	root, proofSet, numLeaves, err := BuildReaderProof(reader, sha256.New(), 2, 0)
	if err != nil {
		t.Fatal(err)
	}

	expectedRoot := sum(sha256.New(), []byte{0, 1})
	if bytes.Compare(root, expectedRoot) != 0 {
		t.Error("ReaderRoot returned the wrong root")
	}
	if len(proofSet) != 1 {
		t.Fatal("proofSet is the incorrect length")
	}
	if bytes.Compare(proofSet[0], []byte{1}) != 0 {
		t.Error("proofSet is incorrect")
	}
	if numLeaves != 1 {
		t.Error("wrong number of leaves returned")
	}
}

// TestEmptyReader passes an empty reader into BuildReaderProof.
func TestEmptyReader(t *testing.T) {
	_, _, _, err := BuildReaderProof(new(bytes.Reader), sha256.New(), 64, 5)
	if err == nil {
		t.Error(err)
	}
}
<|endoftext|>


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

package jujutest

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"sort"
	"time"

	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/bootstrap"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/configstore"
	"launchpad.net/juju-core/environs/network"
	"launchpad.net/juju-core/environs/storage"
	envtesting "launchpad.net/juju-core/environs/testing"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/provider/dummy"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

// Tests is a gocheck suite containing tests verifying juju functionality
// against the environment with the given configuration. The
// tests are not designed to be run against a live server - the Environ
// is opened once for each test, and some potentially expensive operations
// may be executed.
type Tests struct {
	testbase.LoggingSuite
	TestConfig coretesting.Attrs
	envtesting.ToolsFixture

	// ConfigStore holds the configuration storage
	// used when preparing the environment.
	// This is initialized by SetUpTest.
	ConfigStore configstore.Storage
}

// Open opens an instance of the testing environment.
func (t *Tests) Open(c *gc.C) environs.Environ {
	info, err := t.ConfigStore.ReadInfo(t.TestConfig["name"].(string))
	c.Assert(err, gc.IsNil)
	cfg, err := config.New(config.NoDefaults, info.BootstrapConfig())
	c.Assert(err, gc.IsNil)
	e, err := environs.New(cfg)
	c.Assert(err, gc.IsNil, gc.Commentf("opening environ %#v", cfg.AllAttrs()))
	c.Assert(e, gc.NotNil)
	return e
}

// Prepare prepares an instance of the testing environment.
func (t *Tests) Prepare(c *gc.C) environs.Environ {
	cfg, err := config.New(config.NoDefaults, t.TestConfig)
	c.Assert(err, gc.IsNil)
	e, err := environs.Prepare(cfg, coretesting.Context(c), t.ConfigStore)
	c.Assert(err, gc.IsNil, gc.Commentf("preparing environ %#v", t.TestConfig))
	c.Assert(e, gc.NotNil)
	return e
}

func (t *Tests) SetUpTest(c *gc.C) {
	t.LoggingSuite.SetUpTest(c)
	t.ToolsFixture.SetUpTest(c)
	t.ConfigStore = configstore.NewMem()
}

func (t *Tests) TearDownTest(c *gc.C) {
	t.ToolsFixture.TearDownTest(c)
	t.LoggingSuite.TearDownTest(c)
}

func (t *Tests) TestStartStop(c *gc.C) {
	e := t.Prepare(c)
	envtesting.UploadFakeTools(c, e.Storage())
	cfg, err := e.Config().Apply(map[string]interface{}{
		"agent-version": version.Current.Number.String(),
	})
	c.Assert(err, gc.IsNil)
	err = e.SetConfig(cfg)
	c.Assert(err, gc.IsNil)

	insts, err := e.Instances(nil)
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 0)

	inst0, hc := testing.AssertStartInstance(c, e, "0")
	c.Assert(inst0, gc.NotNil)
	id0 := inst0.Id()
	// Sanity check for hardware characteristics.
	c.Assert(hc.Arch, gc.NotNil)
	c.Assert(hc.Mem, gc.NotNil)
	c.Assert(hc.CpuCores, gc.NotNil)

	inst1, _ := testing.AssertStartInstance(c, e, "1")
	c.Assert(inst1, gc.NotNil)
	id1 := inst1.Id()

	insts, err = e.Instances([]instance.Id{id0, id1})
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 2)
	c.Assert(insts[0].Id(), gc.Equals, id0)
	c.Assert(insts[1].Id(), gc.Equals, id1)

	// order of results is not specified
	insts, err = e.AllInstances()
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 2)
	c.Assert(insts[0].Id(), gc.Not(gc.Equals), insts[1].Id())

	err = e.StopInstances([]instance.Instance{inst0})
	c.Assert(err, gc.IsNil)

	insts, err = e.Instances([]instance.Id{id0, id1})
	c.Assert(err, gc.Equals, environs.ErrPartialInstances)
	c.Assert(insts[0], gc.IsNil)
	c.Assert(insts[1].Id(), gc.Equals, id1)

	insts, err = e.AllInstances()
	c.Assert(err, gc.IsNil)
	c.Assert(insts[0].Id(), gc.Equals, id1)
}

func (t *Tests) TestBootstrap(c *gc.C) {
	e := t.prepareAndBootstrap(c)

	info, apiInfo, err := e.StateInfo()
	c.Check(info.Addrs, gc.Not(gc.HasLen), 0)
	c.Check(apiInfo.Addrs, gc.Not(gc.HasLen), 0)

	err = bootstrap.EnsureNotBootstrapped(e)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")

	e2 := t.Open(c)
	envtesting.UploadFakeTools(c, e2.Storage())
	err = bootstrap.EnsureNotBootstrapped(e2)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")

	info2, apiInfo2, err := e2.StateInfo()
	c.Check(info2, gc.DeepEquals, info)
	c.Check(apiInfo2, gc.DeepEquals, apiInfo)

	err = environs.Destroy(e2, t.ConfigStore)
	c.Assert(err, gc.IsNil)

	// Prepare again because Destroy invalidates old environments.
	e3 := t.Prepare(c)
	envtesting.UploadFakeTools(c, e3.Storage())

	err = bootstrap.EnsureNotBootstrapped(e3)
	c.Assert(err, gc.IsNil)
	err = bootstrap.Bootstrap(coretesting.Context(c), e3, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)

	err = bootstrap.EnsureNotBootstrapped(e3)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")
}

func (t *Tests) TestAllocateAddress(c *gc.C) {
	e := t.prepareAndBootstrap(c)
	inst, _ := testing.AssertStartInstance(c, e, "0")
	c.Assert(inst, gc.NotNil)

	opc := make(chan dummy.Operation, 200)
	dummy.Listen(opc)

	expectAddress := instance.NewAddress("0.1.2.1", instance.NetworkCloudLocal)
	assertAddressAllocated(c, e, opc, inst.Id(), network.Id("net1"), expectAddress)

	expectAddress = instance.NewAddress("0.1.2.2", instance.NetworkCloudLocal)
	assertAddressAllocated(c, e, opc, inst.Id(), network.Id("net1"), expectAddress)
}

func (t *Tests) prepareAndBootstrap(c *gc.C) environs.Environ {
	e := t.Prepare(c)
	envtesting.UploadFakeTools(c, e.Storage())
	err := bootstrap.EnsureNotBootstrapped(e)
	c.Assert(err, gc.IsNil)
	err = bootstrap.Bootstrap(coretesting.Context(c), e, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)
	return e
}

func assertAddressAllocated(c *gc.C, e environs.Environ, opc chan dummy.Operation, expectInstId instance.Id, expectNetId network.Id, expectAddress instance.Address) {
	address, err := e.AllocateAddress(expectInstId, expectNetId)
	c.Assert(err, gc.IsNil)
	c.Assert(address, gc.DeepEquals, expectAddress)
	for {
		select {
		case op := <-opc:
			addrOp, ok := op.(dummy.OpAllocateAddress)
			if !ok {
				c.Fatalf("unexpected op: %#v", op)
			}
			c.Check(addrOp.NetworkId, gc.Equals, expectNetId)
			c.Check(addrOp.InstanceId, gc.Equals, expectInstId)
			c.Check(addrOp.Address, gc.Equals, expectAddress)
			return
		case <-time.After(15 * time.Second):
			c.Fatalf("time out wating for operation")
		}
	}
}

var noRetry = utils.AttemptStrategy{}

func (t *Tests) TestPersistence(c *gc.C) {
	stor := t.Prepare(c).Storage()

	names := []string{
		"aa",
		"zzz/aa",
		"zzz/bb",
	}
	for _, name := range names {
		checkFileDoesNotExist(c, stor, name, noRetry)
		checkPutFile(c, stor, name, []byte(name))
	}
	checkList(c, stor, "", names)
	checkList(c, stor, "a", []string{"aa"})
	checkList(c, stor, "zzz/", []string{"zzz/aa", "zzz/bb"})

	storage2 := t.Open(c).Storage()
	for _, name := range names {
		checkFileHasContents(c, storage2, name, []byte(name), noRetry)
	}

	// remove the first file and check that the others remain.
	err := storage2.Remove(names[0])
	c.Check(err, gc.IsNil)

	// check that it's ok to remove a file twice.
	err = storage2.Remove(names[0])
	c.Check(err, gc.IsNil)

	// ... and check it's been removed in the other environment
	checkFileDoesNotExist(c, stor, names[0], noRetry)

	// ... and that the rest of the files are still around
	checkList(c, storage2, "", names[1:])

	for _, name := range names[1:] {
		err := storage2.Remove(name)
		c.Assert(err, gc.IsNil)
	}

	// check they've all gone
	checkList(c, storage2, "", nil)
}

func checkList(c *gc.C, stor storage.StorageReader, prefix string, names []string) {
	lnames, err := storage.List(stor, prefix)
	c.Assert(err, gc.IsNil)
	// TODO(dfc) gocheck should grow an SliceEquals checker.
	expected := copyslice(lnames)
	sort.Strings(expected)
	actual := copyslice(names)
	sort.Strings(actual)
	c.Assert(expected, gc.DeepEquals, actual)
}

// copyslice returns a copy of the slice
func copyslice(s []string) []string {
	r := make([]string, len(s))
	copy(r, s)
	return r
}

func checkPutFile(c *gc.C, stor storage.StorageWriter, name string, contents []byte) {
	err := stor.Put(name, bytes.NewBuffer(contents), int64(len(contents)))
	c.Assert(err, gc.IsNil)
}

func checkFileDoesNotExist(c *gc.C, stor storage.StorageReader, name string, attempt utils.AttemptStrategy) {
	r, err := storage.GetWithRetry(stor, name, attempt)
	c.Assert(r, gc.IsNil)
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func checkFileHasContents(c *gc.C, stor storage.StorageReader, name string, contents []byte, attempt utils.AttemptStrategy) {
	r, err := storage.GetWithRetry(stor, name, attempt)
	c.Assert(err, gc.IsNil)
	c.Check(r, gc.NotNil)
	defer r.Close()

	data, err := ioutil.ReadAll(r)
	c.Check(err, gc.IsNil)
	c.Check(data, gc.DeepEquals, contents)

	url, err := stor.URL(name)
	c.Assert(err, gc.IsNil)

	var resp *http.Response
	for a := attempt.Start(); a.Next(); {
		resp, err = utils.GetValidatingHTTPClient().Get(url)
		c.Assert(err, gc.IsNil)
		if resp.StatusCode != 404 {
			break
		}
		c.Logf("get retrying after earlier get succeeded. *sigh*.")
	}
	c.Assert(err, gc.IsNil)
	data, err = ioutil.ReadAll(resp.Body)
	c.Assert(err, gc.IsNil)
	defer resp.Body.Close()
	c.Assert(resp.StatusCode, gc.Equals, 200, gc.Commentf("error response: %s", data))
	c.Check(data, gc.DeepEquals, contents)
}
<commit_msg>Changes after review<commit_after>// Copyright 2011, 2012, 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package jujutest

import (
	"bytes"
	"io/ioutil"
	"net/http"
	"sort"
	"time"

	jc "github.com/juju/testing/checkers"
	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/environs/bootstrap"
	"launchpad.net/juju-core/environs/config"
	"launchpad.net/juju-core/environs/configstore"
	"launchpad.net/juju-core/environs/network"
	"launchpad.net/juju-core/environs/storage"
	envtesting "launchpad.net/juju-core/environs/testing"
	"launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/provider/dummy"
	coretesting "launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/testing/testbase"
	"launchpad.net/juju-core/utils"
	"launchpad.net/juju-core/version"
)

// Tests is a gocheck suite containing tests verifying juju functionality
// against the environment with the given configuration. The
// tests are not designed to be run against a live server - the Environ
// is opened once for each test, and some potentially expensive operations
// may be executed.
type Tests struct {
	testbase.LoggingSuite
	TestConfig coretesting.Attrs
	envtesting.ToolsFixture

	// ConfigStore holds the configuration storage
	// used when preparing the environment.
	// This is initialized by SetUpTest.
	ConfigStore configstore.Storage
}

// Open opens an instance of the testing environment.
func (t *Tests) Open(c *gc.C) environs.Environ {
	info, err := t.ConfigStore.ReadInfo(t.TestConfig["name"].(string))
	c.Assert(err, gc.IsNil)
	cfg, err := config.New(config.NoDefaults, info.BootstrapConfig())
	c.Assert(err, gc.IsNil)
	e, err := environs.New(cfg)
	c.Assert(err, gc.IsNil, gc.Commentf("opening environ %#v", cfg.AllAttrs()))
	c.Assert(e, gc.NotNil)
	return e
}

// Prepare prepares an instance of the testing environment.
func (t *Tests) Prepare(c *gc.C) environs.Environ {
	cfg, err := config.New(config.NoDefaults, t.TestConfig)
	c.Assert(err, gc.IsNil)
	e, err := environs.Prepare(cfg, coretesting.Context(c), t.ConfigStore)
	c.Assert(err, gc.IsNil, gc.Commentf("preparing environ %#v", t.TestConfig))
	c.Assert(e, gc.NotNil)
	return e
}

func (t *Tests) SetUpTest(c *gc.C) {
	t.LoggingSuite.SetUpTest(c)
	t.ToolsFixture.SetUpTest(c)
	t.ConfigStore = configstore.NewMem()
}

func (t *Tests) TearDownTest(c *gc.C) {
	t.ToolsFixture.TearDownTest(c)
	t.LoggingSuite.TearDownTest(c)
}

func (t *Tests) TestStartStop(c *gc.C) {
	e := t.Prepare(c)
	envtesting.UploadFakeTools(c, e.Storage())
	cfg, err := e.Config().Apply(map[string]interface{}{
		"agent-version": version.Current.Number.String(),
	})
	c.Assert(err, gc.IsNil)
	err = e.SetConfig(cfg)
	c.Assert(err, gc.IsNil)

	insts, err := e.Instances(nil)
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 0)

	inst0, hc := testing.AssertStartInstance(c, e, "0")
	c.Assert(inst0, gc.NotNil)
	id0 := inst0.Id()
	// Sanity check for hardware characteristics.
	c.Assert(hc.Arch, gc.NotNil)
	c.Assert(hc.Mem, gc.NotNil)
	c.Assert(hc.CpuCores, gc.NotNil)

	inst1, _ := testing.AssertStartInstance(c, e, "1")
	c.Assert(inst1, gc.NotNil)
	id1 := inst1.Id()

	insts, err = e.Instances([]instance.Id{id0, id1})
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 2)
	c.Assert(insts[0].Id(), gc.Equals, id0)
	c.Assert(insts[1].Id(), gc.Equals, id1)

	// order of results is not specified
	insts, err = e.AllInstances()
	c.Assert(err, gc.IsNil)
	c.Assert(insts, gc.HasLen, 2)
	c.Assert(insts[0].Id(), gc.Not(gc.Equals), insts[1].Id())

	err = e.StopInstances([]instance.Instance{inst0})
	c.Assert(err, gc.IsNil)

	insts, err = e.Instances([]instance.Id{id0, id1})
	c.Assert(err, gc.Equals, environs.ErrPartialInstances)
	c.Assert(insts[0], gc.IsNil)
	c.Assert(insts[1].Id(), gc.Equals, id1)

	insts, err = e.AllInstances()
	c.Assert(err, gc.IsNil)
	c.Assert(insts[0].Id(), gc.Equals, id1)
}

func (t *Tests) TestBootstrap(c *gc.C) {
	e := t.prepareAndBootstrap(c)

	info, apiInfo, err := e.StateInfo()
	c.Check(info.Addrs, gc.Not(gc.HasLen), 0)
	c.Check(apiInfo.Addrs, gc.Not(gc.HasLen), 0)

	err = bootstrap.EnsureNotBootstrapped(e)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")

	e2 := t.Open(c)
	envtesting.UploadFakeTools(c, e2.Storage())
	err = bootstrap.EnsureNotBootstrapped(e2)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")

	info2, apiInfo2, err := e2.StateInfo()
	c.Check(info2, gc.DeepEquals, info)
	c.Check(apiInfo2, gc.DeepEquals, apiInfo)

	err = environs.Destroy(e2, t.ConfigStore)
	c.Assert(err, gc.IsNil)

	// Prepare again because Destroy invalidates old environments.
	e3 := t.Prepare(c)
	envtesting.UploadFakeTools(c, e3.Storage())

	err = bootstrap.EnsureNotBootstrapped(e3)
	c.Assert(err, gc.IsNil)
	err = bootstrap.Bootstrap(coretesting.Context(c), e3, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)

	err = bootstrap.EnsureNotBootstrapped(e3)
	c.Assert(err, gc.ErrorMatches, "environment is already bootstrapped")
}

func (t *Tests) TestAllocateAddress(c *gc.C) {
	e := t.prepareAndBootstrap(c)
	inst, _ := testing.AssertStartInstance(c, e, "0")
	c.Assert(inst, gc.NotNil)
	netId := network.Id("net1")

	opc := make(chan dummy.Operation, 200)
	dummy.Listen(opc)

	expectAddress := instance.NewAddress("0.1.2.1", instance.NetworkCloudLocal)
	address, err := e.AllocateAddress(inst.Id(), netId)
	c.Assert(err, gc.IsNil)
	c.Assert(address, gc.DeepEquals, expectAddress)

	assertAllocateAddress(c, e, opc, inst.Id(), netId, expectAddress)

	expectAddress = instance.NewAddress("0.1.2.2", instance.NetworkCloudLocal)
	address, err = e.AllocateAddress(inst.Id(), netId)
	c.Assert(err, gc.IsNil)
	c.Assert(address, gc.DeepEquals, expectAddress)
	assertAllocateAddress(c, e, opc, inst.Id(), netId, expectAddress)
}

func (t *Tests) prepareAndBootstrap(c *gc.C) environs.Environ {
	e := t.Prepare(c)
	envtesting.UploadFakeTools(c, e.Storage())
	err := bootstrap.EnsureNotBootstrapped(e)
	c.Assert(err, gc.IsNil)
	err = bootstrap.Bootstrap(coretesting.Context(c), e, environs.BootstrapParams{})
	c.Assert(err, gc.IsNil)
	return e
}

func assertAllocateAddress(c *gc.C, e environs.Environ, opc chan dummy.Operation, expectInstId instance.Id, expectNetId network.Id, expectAddress instance.Address) {
	select {
	case op := <-opc:
		addrOp, ok := op.(dummy.OpAllocateAddress)
		if !ok {
			c.Fatalf("unexpected op: %#v", op)
		}
		c.Check(addrOp.NetworkId, gc.Equals, expectNetId)
		c.Check(addrOp.InstanceId, gc.Equals, expectInstId)
		c.Check(addrOp.Address, gc.Equals, expectAddress)
		return
	case <-time.After(coretesting.ShortWait):
		c.Fatalf("time out wating for operation")
	}
}

var noRetry = utils.AttemptStrategy{}

func (t *Tests) TestPersistence(c *gc.C) {
	stor := t.Prepare(c).Storage()

	names := []string{
		"aa",
		"zzz/aa",
		"zzz/bb",
	}
	for _, name := range names {
		checkFileDoesNotExist(c, stor, name, noRetry)
		checkPutFile(c, stor, name, []byte(name))
	}
	checkList(c, stor, "", names)
	checkList(c, stor, "a", []string{"aa"})
	checkList(c, stor, "zzz/", []string{"zzz/aa", "zzz/bb"})

	storage2 := t.Open(c).Storage()
	for _, name := range names {
		checkFileHasContents(c, storage2, name, []byte(name), noRetry)
	}

	// remove the first file and check that the others remain.
	err := storage2.Remove(names[0])
	c.Check(err, gc.IsNil)

	// check that it's ok to remove a file twice.
	err = storage2.Remove(names[0])
	c.Check(err, gc.IsNil)

	// ... and check it's been removed in the other environment
	checkFileDoesNotExist(c, stor, names[0], noRetry)

	// ... and that the rest of the files are still around
	checkList(c, storage2, "", names[1:])

	for _, name := range names[1:] {
		err := storage2.Remove(name)
		c.Assert(err, gc.IsNil)
	}

	// check they've all gone
	checkList(c, storage2, "", nil)
}

func checkList(c *gc.C, stor storage.StorageReader, prefix string, names []string) {
	lnames, err := storage.List(stor, prefix)
	c.Assert(err, gc.IsNil)
	// TODO(dfc) gocheck should grow an SliceEquals checker.
	expected := copyslice(lnames)
	sort.Strings(expected)
	actual := copyslice(names)
	sort.Strings(actual)
	c.Assert(expected, gc.DeepEquals, actual)
}

// copyslice returns a copy of the slice
func copyslice(s []string) []string {
	r := make([]string, len(s))
	copy(r, s)
	return r
}

func checkPutFile(c *gc.C, stor storage.StorageWriter, name string, contents []byte) {
	err := stor.Put(name, bytes.NewBuffer(contents), int64(len(contents)))
	c.Assert(err, gc.IsNil)
}

func checkFileDoesNotExist(c *gc.C, stor storage.StorageReader, name string, attempt utils.AttemptStrategy) {
	r, err := storage.GetWithRetry(stor, name, attempt)
	c.Assert(r, gc.IsNil)
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
}

func checkFileHasContents(c *gc.C, stor storage.StorageReader, name string, contents []byte, attempt utils.AttemptStrategy) {
	r, err := storage.GetWithRetry(stor, name, attempt)
	c.Assert(err, gc.IsNil)
	c.Check(r, gc.NotNil)
	defer r.Close()

	data, err := ioutil.ReadAll(r)
	c.Check(err, gc.IsNil)
	c.Check(data, gc.DeepEquals, contents)

	url, err := stor.URL(name)
	c.Assert(err, gc.IsNil)

	var resp *http.Response
	for a := attempt.Start(); a.Next(); {
		resp, err = utils.GetValidatingHTTPClient().Get(url)
		c.Assert(err, gc.IsNil)
		if resp.StatusCode != 404 {
			break
		}
		c.Logf("get retrying after earlier get succeeded. *sigh*.")
	}
	c.Assert(err, gc.IsNil)
	data, err = ioutil.ReadAll(resp.Body)
	c.Assert(err, gc.IsNil)
	defer resp.Body.Close()
	c.Assert(resp.StatusCode, gc.Equals, 200, gc.Commentf("error response: %s", data))
	c.Check(data, gc.DeepEquals, contents)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-zigbee/gateway"
)

type OnOffChannel struct {
	Channel
	channel *channels.OnOffChannel
}

// -------- On/Off Protocol --------

func (c *OnOffChannel) TurnOn() error {
	return c.setState(gateway.GwOnOffStateT_ON_STATE.Enum())
}

func (c *OnOffChannel) TurnOff() error {
	return c.setState(gateway.GwOnOffStateT_OFF_STATE.Enum())
}

func (c *OnOffChannel) ToggleOnOff() error {
	return c.setState(gateway.GwOnOffStateT_TOGGLE_STATE.Enum())
}

func (c *OnOffChannel) SetOnOff(state bool) error {
	if state {
		return c.TurnOn()
	}

	return c.TurnOff()
}

func (c *OnOffChannel) init() error {
	log.Printf("Initialising on/off channel of device %d", *c.device.deviceInfo.IeeeAddress)

	clusterID := uint32(0x06)
	attributeID := uint32(0)
	minReportInterval := uint32(1)
	maxReportInterval := uint32(120)

	request := &gateway.GwSetAttributeReportingReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		ClusterId: &clusterID,
		AttributeReportList: []*gateway.GwAttributeReportT{{
			AttributeId:       &attributeID,
			AttributeType:     gateway.GwZclAttributeDataTypesT_ZCL_DATATYPE_BOOLEAN.Enum(),
			MinReportInterval: &minReportInterval,
			MaxReportInterval: &maxReportInterval,
		}},
	}

	response := &gateway.GwSetAttributeReportingRspInd{}

	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 20*time.Second)
	if err != nil {
		log.Printf("Error enabling on/off reporting: %s", err)
	} else if response.Status.String() != "STATUS_SUCCESS" {
		log.Printf("Failed to enable on/off reporting. status: %s", response.Status.String())
	}

	c.channel = channels.NewOnOffChannel(c)
	err = c.device.driver.conn.ExportChannel(c.device, c.channel, "on-off")
	if err != nil {
		log.Fatalf("Failed to announce on/off channel: %s", err)
	}

	go func() {
		for {
			log.Printf("Polling for on/off")
			err := c.fetchState()
			if err != nil {
				log.Printf("Failed to poll for on/off state %s", err)
			}
			time.Sleep(10 * time.Second)
		}
	}()

	return nil

}

func (c *OnOffChannel) setState(state *gateway.GwOnOffStateT) error {
	request := &gateway.DevSetOnOffStateReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		State: gateway.GwOnOffStateT_TOGGLE_STATE.Enum(),
	}

	response := &gateway.GwZigbeeGenericRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 2*time.Second)
	if err != nil {
		return fmt.Errorf("Error setting on/off state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to set on/off state. status: %s", response.Status.String())
	}

	return c.fetchState()
}

func (c *OnOffChannel) fetchState() error {
	request := &gateway.DevGetOnOffStateReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
	}

	response := &gateway.DevGetOnOffStateRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 10*time.Second)
	if err != nil {
		return fmt.Errorf("Error getting on/off state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to get on/off state. status: %s", response.Status.String())
	}

	c.channel.SendState(*response.StateValue == gateway.GwOnOffStateValueT_ON)

	return nil
}
<commit_msg>Use requested state to implement on/off.<commit_after>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/ninjasphere/go-ninja/channels"
	"github.com/ninjasphere/go-zigbee/gateway"
)

type OnOffChannel struct {
	Channel
	channel *channels.OnOffChannel
}

// -------- On/Off Protocol --------

func (c *OnOffChannel) TurnOn() error {
	return c.setState(gateway.GwOnOffStateT_ON_STATE.Enum())
}

func (c *OnOffChannel) TurnOff() error {
	return c.setState(gateway.GwOnOffStateT_OFF_STATE.Enum())
}

func (c *OnOffChannel) ToggleOnOff() error {
	return c.setState(gateway.GwOnOffStateT_TOGGLE_STATE.Enum())
}

func (c *OnOffChannel) SetOnOff(state bool) error {
	if state {
		return c.TurnOn()
	}

	return c.TurnOff()
}

func (c *OnOffChannel) init() error {
	log.Printf("Initialising on/off channel of device %d", *c.device.deviceInfo.IeeeAddress)

	clusterID := uint32(0x06)
	attributeID := uint32(0)
	minReportInterval := uint32(1)
	maxReportInterval := uint32(120)

	request := &gateway.GwSetAttributeReportingReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		ClusterId: &clusterID,
		AttributeReportList: []*gateway.GwAttributeReportT{{
			AttributeId:       &attributeID,
			AttributeType:     gateway.GwZclAttributeDataTypesT_ZCL_DATATYPE_BOOLEAN.Enum(),
			MinReportInterval: &minReportInterval,
			MaxReportInterval: &maxReportInterval,
		}},
	}

	response := &gateway.GwSetAttributeReportingRspInd{}

	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 20*time.Second)
	if err != nil {
		log.Printf("Error enabling on/off reporting: %s", err)
	} else if response.Status.String() != "STATUS_SUCCESS" {
		log.Printf("Failed to enable on/off reporting. status: %s", response.Status.String())
	}

	c.channel = channels.NewOnOffChannel(c)
	err = c.device.driver.conn.ExportChannel(c.device, c.channel, "on-off")
	if err != nil {
		log.Fatalf("Failed to announce on/off channel: %s", err)
	}

	go func() {
		for {
			log.Printf("Polling for on/off")
			err := c.fetchState()
			if err != nil {
				log.Printf("Failed to poll for on/off state %s", err)
			}
			time.Sleep(10 * time.Second)
		}
	}()

	return nil

}

func (c *OnOffChannel) setState(state *gateway.GwOnOffStateT) error {
	request := &gateway.DevSetOnOffStateReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
		State: state,
	}

	response := &gateway.GwZigbeeGenericRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 2*time.Second)
	if err != nil {
		return fmt.Errorf("Error setting on/off state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to set on/off state. status: %s", response.Status.String())
	}

	return c.fetchState()
}

func (c *OnOffChannel) fetchState() error {
	request := &gateway.DevGetOnOffStateReq{
		DstAddress: &gateway.GwAddressStructT{
			AddressType: gateway.GwAddressTypeT_UNICAST.Enum(),
			IeeeAddr:    c.device.deviceInfo.IeeeAddress,
		},
	}

	response := &gateway.DevGetOnOffStateRspInd{}
	err := c.device.driver.gatewayConn.SendAsyncCommand(request, response, 10*time.Second)
	if err != nil {
		return fmt.Errorf("Error getting on/off state : %s", err)
	}
	if response.Status.String() != "STATUS_SUCCESS" {
		return fmt.Errorf("Failed to get on/off state. status: %s", response.Status.String())
	}

	c.channel.SendState(*response.StateValue == gateway.GwOnOffStateValueT_ON)

	return nil
}
<|endoftext|>


<commit_before>package jwthelper_test

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/northbright/jwthelper"
)

var (
	serverURL = "localhost:8080"
)

func doPostRequest(URL string) {
	v := url.Values{}
	v.Set("username", "admin")
	v.Set("password", "admin")

	// Values.Encode() encodes the values into "URL encoded" form sorted by key.
	s := v.Encode()

	req, err := http.NewRequest("POST", URL, strings.NewReader(s))
	if err != nil {
		log.Printf("NewRequest error: %v", err)
		return
	}

	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	c := &http.Client{}
	resp, err := c.Do(req)
	if err != nil {
		log.Printf("Do() error: %v", err)
		return
	}
	defer resp.Body.Close()

	// Get Cookies
	cookies := resp.Cookies()
	log.Printf("After POST, cookies: %v, resp: %v", cookies, resp)
}

func login(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		fmt.Fprintf(w, "GET")
	case "POST":
		// Call ParseForm() to parse the raw query and update r.PostForm and r.Form.
		if err := r.ParseForm(); err != nil {
			return
		}

		// Post form from website
		username := r.FormValue("username")
		password := r.FormValue("password")
		if username == "admin" && password == "admin" {
			signer := jwthelper.NewRSASHASigner([]byte(rsaPrivPEM))
			tokenString, err := signer.SignedString(
				jwthelper.NewClaim("username", username),
			)
			if err != nil {
				return
			}
			cookie := jwthelper.NewCookie(tokenString)
			http.SetCookie(w, cookie)
			fmt.Fprintf(w, "POST")
		}
	default:
		fmt.Fprintf(w, "Sorry, only GET and POST methods are supported.")
	}
}

func shutdownServer(srv *http.Server) {
	log.Printf("shutdown server...")
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := srv.Shutdown(ctx); err != nil {
		log.Printf("shutdown server error: %v", err)
	}
	log.Println("shutdown server successfully")
}

func ExampleNewCookie() {
	mux := http.NewServeMux()
	mux.HandleFunc("/login", login)

	srv := &http.Server{
		Addr:    serverURL,
		Handler: mux,
	}

	go func() {
		time.Sleep(1 * time.Second)
		doPostRequest("http://" + serverURL + "/login")
		shutdownServer(srv)
	}()

	err := srv.ListenAndServe()
	if err != nil {
		if err == http.ErrServerClosed {
			log.Printf("server has been closed")
		} else {
			log.Printf("ListenAndServe() error: %s", err)
		}
	}

	// Output:
}
<commit_msg>Add doGetRequest() to test get JWT cookie in request<commit_after>package jwthelper_test

import (
	"context"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/northbright/jwthelper"
)

func doPostRequest(URL string) *http.Cookie {
	v := url.Values{}
	v.Set("username", "admin")
	v.Set("password", "admin")

	// Values.Encode() encodes the values into "URL encoded" form sorted by key.
	s := v.Encode()

	req, err := http.NewRequest("POST", URL, strings.NewReader(s))
	if err != nil {
		log.Printf("NewRequest error: %v", err)
		return nil
	}

	req.Header.Add("Content-Type", "application/x-www-form-urlencoded")

	c := &http.Client{}
	resp, err := c.Do(req)
	if err != nil {
		log.Printf("Do() error: %v", err)
		return nil
	}
	defer resp.Body.Close()

	// Get JWT cookie("jwt").
	cookies := resp.Cookies()
	for _, cookie := range cookies {
		if cookie.Name == "jwt" {
			log.Printf("After POST, JWT cookie: %v, resp: %v", cookie, resp)
			return cookie
		}
	}
	return nil
}

func doGetRequest(URL string, cookie *http.Cookie) {
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		log.Printf("NewRequest error: %v", err)
		return
	}
	// Add JWT cookie return by POST.
	req.AddCookie(cookie)

	c := &http.Client{}
	resp, err := c.Do(req)
	if err != nil {
		log.Printf("Do() error: %v", err)
		return
	}
	defer resp.Body.Close()

	// Get response("admin").
	buf, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Printf("ReadAll() error: %v", err)
		return
	}
	log.Printf("GET response: %v", string(buf))
}

func login(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		cookie, err := r.Cookie("jwt")
		if err != nil {
			log.Printf("get JWT cookie error: %v", err)
			return
		}

		tokenString := cookie.Value
		parser := jwthelper.NewRSASHAParser([]byte(rsaPubPEM))
		m, err := parser.Parse(tokenString)
		if err != nil {
			log.Printf("parser.Parse() error: %v", err)
			return
		}
		fmt.Fprintf(w, "hello, %v!", m["username"])

	case "POST":
		// Call ParseForm() to parse the raw query and update r.PostForm and r.Form.
		if err := r.ParseForm(); err != nil {
			return
		}

		// Post form from website
		username := r.FormValue("username")
		password := r.FormValue("password")
		if username == "admin" && password == "admin" {
			signer := jwthelper.NewRSASHASigner([]byte(rsaPrivPEM))
			tokenString, err := signer.SignedString(
				jwthelper.NewClaim("username", username),
			)
			if err != nil {
				return
			}
			cookie := jwthelper.NewCookie(tokenString)
			http.SetCookie(w, cookie)
			fmt.Fprintf(w, "POST")
		}
	default:
		fmt.Fprintf(w, "Sorry, only GET and POST methods are supported.")
	}
}

func shutdownServer(srv *http.Server) {
	log.Printf("shutdown server...")
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := srv.Shutdown(ctx); err != nil {
		log.Printf("shutdown server error: %v", err)
	}
	log.Println("shutdown server successfully")
}

func ExampleNewCookie() {
	mux := http.NewServeMux()
	mux.HandleFunc("/login", login)

	srv := &http.Server{
		Addr:    ":8080",
		Handler: mux,
	}

	go func() {
		time.Sleep(1 * time.Second)
		cookie := doPostRequest("http://localhost:8080/login")
		if cookie != nil {
			doGetRequest("http://localhost:8080/login", cookie)
		}
		shutdownServer(srv)
	}()

	err := srv.ListenAndServe()
	if err != nil {
		if err == http.ErrServerClosed {
			log.Printf("server has been closed")
		} else {
			log.Printf("ListenAndServe() error: %s", err)
		}
	}

	// Output:
}
<|endoftext|>