<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 obcpbft

import (
	"encoding/base64"
	"fmt"
	"time"

	"github.com/openblockchain/obc-peer/openchain/consensus"
	"github.com/openblockchain/obc-peer/openchain/util"
	pb "github.com/openblockchain/obc-peer/protos"

	"github.com/golang/protobuf/proto"
	"github.com/spf13/viper"
)

type obcBatch struct {
	cpi  consensus.CPI
	pbft *pbftCore

	batchSize        int
	batchStore       map[string]*Request
	batchTimer       *time.Timer
	batchTimerActive bool
	batchTimeout     time.Duration
}

func newObcBatch(id uint64, config *viper.Viper, cpi consensus.CPI) *obcBatch {
	var err error
	op := &obcBatch{cpi: cpi}
	op.pbft = newPbftCore(id, config, op, cpi)
	op.batchSize = config.GetInt("general.batchSize")
	op.batchStore = make(map[string]*Request)
	op.batchTimeout, err = time.ParseDuration(config.GetString("general.timeout.batch"))
	if err != nil {
		panic(fmt.Errorf("Cannot parse batch timeout: %s", err))
	}
	// create non-running timer XXX ugly
	op.batchTimer = time.NewTimer(100 * time.Hour)
	op.batchTimer.Stop()
	go op.batchTimerHander()
	return op
}

// RecvMsg receives both CHAIN_TRANSACTION and CONSENSUS messages from
// the stack. New transaction requests are broadcast to all replicas,
// so that the current primary will receive the request.
func (op *obcBatch) RecvMsg(ocMsg *pb.OpenchainMessage) error {
	if ocMsg.Type == pb.OpenchainMessage_CHAIN_TRANSACTION {
		logger.Info("New consensus request received")

		if err := op.validate(ocMsg.Payload); err != nil {
			logger.Warning("Request did not verify: %s", err)
			return err
		}

		req := &Request{Payload: ocMsg.Payload, ReplicaId: op.pbft.id}

		if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView { // primary
			err := op.leaderProcReq(req)
			if err != nil {
				return err
			}
		} else { // backup
			msg := &Message{&Message_Request{req}}
			msgRaw, _ := proto.Marshal(msg)
			op.broadcast(msgRaw)
		}
		return nil
	}

	if ocMsg.Type != pb.OpenchainMessage_CONSENSUS {
		return fmt.Errorf("Unexpected message type: %s", ocMsg.Type)
	}

	pbftMsg := &Message{}
	err := proto.Unmarshal(ocMsg.Payload, pbftMsg)
	if err != nil {
		return err
	}

	if req := pbftMsg.GetRequest(); req != nil {
		if err = op.validate(req.Payload); err != nil {
			logger.Warning("Request did not verify: %s", err)
			return err
		}

		switch req.ReplicaId {
		case op.pbft.primary(op.pbft.view):
			// a request sent by the primary; primary should ignore this
			if op.pbft.primary(op.pbft.view) != op.pbft.id {
				op.pbft.request(req.Payload)
			}
		default:
			// a request sent by a backup; backups should ignore this
			if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView {
				err := op.leaderProcReq(req)
				if err != nil {
					return err
				}
			}
		}
	} else {
		op.pbft.receive(ocMsg.Payload)
	}

	return nil
}

// Close tells us to release resources we are holding
func (op *obcBatch) Close() {
	op.pbft.close()
	op.batchTimer.Reset(0)
}

// =============================================================================
// innerCPI interface (functions called by pbft-core)
// =============================================================================

// multicast a message to all replicas
func (op *obcBatch) broadcast(msgPayload []byte) {
	ocMsg := &pb.OpenchainMessage{
		Type:    pb.OpenchainMessage_CONSENSUS,
		Payload: msgPayload,
	}
	op.cpi.Broadcast(ocMsg)
}

// send a message to a specific replica
func (op *obcBatch) unicast(msgPayload []byte, receiverID uint64) (err error) {
	ocMsg := &pb.OpenchainMessage{
		Type:    pb.OpenchainMessage_CONSENSUS,
		Payload: msgPayload,
	}
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return
	}
	return op.cpi.Unicast(ocMsg, receiverHandle)
}

func (op *obcBatch) sign(msg []byte) ([]byte, error) {
	return op.cpi.Sign(msg)
}

// verify message signature
func (op *obcBatch) verify(senderID uint64, signature []byte, message []byte) error {
	senderHandle, err := getValidatorHandle(senderID)
	if err != nil {
		return fmt.Errorf("Could not verify message from %v: %v", senderHandle.Name, err)
	}
	return op.cpi.Verify(senderHandle, signature, message)
}

// validate checks whether the request is valid syntactically.
// For now, we only need this for the obc-sieve verify/verify-set and flush messages.
// Thus, for obc-batch, this is a no-op.
func (op *obcBatch) validate(txRaw []byte) error {
	return nil
}

// execute an opaque request which corresponds to an OBC Transaction
func (op *obcBatch) execute(tbRaw []byte, rawMetadata []byte) {
	tb := &pb.TransactionBlock{}
	err := proto.Unmarshal(tbRaw, tb)
	if err != nil {
		return
	}

	txs := tb.Transactions
	txBatchID := base64.StdEncoding.EncodeToString(util.ComputeCryptoHash(tbRaw))

	for i, tx := range txs {
		txRaw, _ := proto.Marshal(tx)
		if err = op.validate(txRaw); err != nil {
			logger.Error("Request in transaction %d from batch %s did not verify: %s", i, txBatchID, err)
			return
		}
	}

	if err := op.cpi.BeginTxBatch(txBatchID); err != nil {
		logger.Error("Failed to begin transaction batch %s: %v", txBatchID, err)
		return
	}

	_, errs := op.cpi.ExecTXs(txs)
	if errs[len(txs)] != nil {
		logger.Error("Fail to execute transaction batch %s: %v", txBatchID, errs)
		if err = op.cpi.RollbackTxBatch(txBatchID); err != nil {
			panic(fmt.Errorf("Unable to rollback transaction batch %s: %v", txBatchID, err))
		}
		return
	}

	if err = op.cpi.CommitTxBatch(txBatchID, txs, nil, rawMetadata); err != nil {
		logger.Error("Failed to commit transaction batch %s to the ledger: %v", txBatchID, err)
		if err = op.cpi.RollbackTxBatch(txBatchID); err != nil {
			panic(fmt.Errorf("Unable to rollback transaction batch %s: %v", txBatchID, err))
		}
		return
	}
}

// signal when a view-change happened
func (op *obcBatch) viewChange(curView uint64) {
	if op.batchTimerActive {
		op.stopBatchTimer()
	}
}

// =============================================================================
// functions specific to batch mode
// =============================================================================

func (op *obcBatch) leaderProcReq(req *Request) error {
	digest := hashReq(req)
	op.batchStore[digest] = req

	if !op.batchTimerActive {
		op.startBatchTimer()
	}

	if len(op.batchStore) >= op.batchSize {
		op.sendBatch()
	}

	return nil
}

func (op *obcBatch) sendBatch() error {
	op.stopBatchTimer()
	// assemble new Request message
	txs := make([]*pb.Transaction, len(op.batchStore))
	var i int
	for d, req := range op.batchStore {
		txs[i] = &pb.Transaction{}
		err := proto.Unmarshal(req.Payload, txs[i])
		if err != nil {
			err = fmt.Errorf("Unable to unpack payload of request %d", i)
			logger.Error("%s", err)
			return err
		}
		i++
		delete(op.batchStore, d) // clean up
	}
	tb := &pb.TransactionBlock{Transactions: txs}
	tbPacked, err := proto.Marshal(tb)
	if err != nil {
		err = fmt.Errorf("Unable to pack transaction block for new batch request")
		logger.Error("%s", err)
		return err
	}
	// process internally
	op.pbft.request(tbPacked)
	// broadcast
	batchReq := &Request{Payload: tbPacked}
	msg := &Message{&Message_Request{batchReq}}
	msgRaw, _ := proto.Marshal(msg)
	op.broadcast(msgRaw)

	return nil
}

// allow the primary to send a batch when the timer expires
func (op *obcBatch) batchTimerHander() {
	for {
		select {
		case <-op.batchTimer.C:
			op.pbft.lock.Lock()
			if op.pbft.closed {
				op.pbft.lock.Unlock()
				return
			}
			logger.Info("Replica %d batch timeout expired", op.pbft.id)
			if op.pbft.activeView && (len(op.batchStore) > 0) {
				op.sendBatch()
			}
			op.pbft.lock.Unlock()
		}
	}
}

func (op *obcBatch) startBatchTimer() {
	op.batchTimer.Reset(op.batchTimeout)
	logger.Debug("Replica %d started the batch timer", op.pbft.id)
	op.batchTimerActive = true
}

func (op *obcBatch) stopBatchTimer() {
	op.batchTimer.Stop()
	logger.Debug("Replica %d stopped the batch timer", op.pbft.id)
	op.batchTimerActive = false
	if op.pbft.closed {
		return
	}
loopBatch:
	for {
		select {
		case <-op.batchTimer.C:
		default:
			break loopBatch
		}
	}
}
<commit_msg>Switch obc-batch to new RecvMsg(msg,senderHandle)<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 obcpbft

import (
	"encoding/base64"
	"fmt"
	"time"

	"github.com/openblockchain/obc-peer/openchain/consensus"
	"github.com/openblockchain/obc-peer/openchain/util"
	pb "github.com/openblockchain/obc-peer/protos"

	"github.com/golang/protobuf/proto"
	"github.com/spf13/viper"
)

type obcBatch struct {
	cpi  consensus.CPI
	pbft *pbftCore

	batchSize        int
	batchStore       map[string]*Request
	batchTimer       *time.Timer
	batchTimerActive bool
	batchTimeout     time.Duration
}

func newObcBatch(id uint64, config *viper.Viper, cpi consensus.CPI) *obcBatch {
	var err error
	op := &obcBatch{cpi: cpi}
	op.pbft = newPbftCore(id, config, op, cpi)
	op.batchSize = config.GetInt("general.batchSize")
	op.batchStore = make(map[string]*Request)
	op.batchTimeout, err = time.ParseDuration(config.GetString("general.timeout.batch"))
	if err != nil {
		panic(fmt.Errorf("Cannot parse batch timeout: %s", err))
	}
	// create non-running timer XXX ugly
	op.batchTimer = time.NewTimer(100 * time.Hour)
	op.batchTimer.Stop()
	go op.batchTimerHander()
	return op
}

// RecvMsg receives both CHAIN_TRANSACTION and CONSENSUS messages from
// the stack. New transaction requests are broadcast to all replicas,
// so that the current primary will receive the request.
func (op *obcBatch) RecvMsg(ocMsg *pb.OpenchainMessage, senderHandle *pb.PeerID) error {
	if ocMsg.Type == pb.OpenchainMessage_CHAIN_TRANSACTION {
		logger.Info("New consensus request received")

		if err := op.validate(ocMsg.Payload); err != nil {
			logger.Warning("Request did not verify: %s", err)
			return err
		}

		req := &Request{Payload: ocMsg.Payload, ReplicaId: op.pbft.id}

		if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView { // primary
			err := op.leaderProcReq(req)
			if err != nil {
				return err
			}
		} else { // backup
			msg := &Message{&Message_Request{req}}
			msgRaw, _ := proto.Marshal(msg)
			op.broadcast(msgRaw)
		}
		return nil
	}

	if ocMsg.Type != pb.OpenchainMessage_CONSENSUS {
		return fmt.Errorf("Unexpected message type: %s", ocMsg.Type)
	}

	senderID, err := getValidatorID(senderHandle)
	if err != nil {
		panic("Cannot map sender's PeerID to a valid replica ID")
	}

	pbftMsg := &Message{}
	err := proto.Unmarshal(ocMsg.Payload, pbftMsg)
	if err != nil {
		return err
	}

	if req := pbftMsg.GetRequest(); req != nil {
		if err = op.validate(req.Payload); err != nil {
			logger.Warning("Request did not verify: %s", err)
			return err
		}

		if senderID != req.ReplicaId {
			err := fmt.Errorf("Sender ID included in message (%v) doesn't match ID corresponding to the receiving stream (%v)", req.ReplicaId, senderID)
			logger.Warning(err.Error())
			return err
		}

		switch req.ReplicaId {
		case op.pbft.primary(op.pbft.view):
			// a request sent by the primary; primary should ignore this
			if op.pbft.primary(op.pbft.view) != op.pbft.id {
				op.pbft.request(req.Payload, senderID)
			}
		default:
			// a request sent by a backup; backups should ignore this
			if (op.pbft.primary(op.pbft.view) == op.pbft.id) && op.pbft.activeView {
				err := op.leaderProcReq(req)
				if err != nil {
					return err
				}
			}
		}
	} else {
		op.pbft.receive(ocMsg.Payload, senderID)
	}

	return nil
}

// Close tells us to release resources we are holding
func (op *obcBatch) Close() {
	op.pbft.close()
	op.batchTimer.Reset(0)
}

// =============================================================================
// innerCPI interface (functions called by pbft-core)
// =============================================================================

// multicast a message to all replicas
func (op *obcBatch) broadcast(msgPayload []byte) {
	ocMsg := &pb.OpenchainMessage{
		Type:    pb.OpenchainMessage_CONSENSUS,
		Payload: msgPayload,
	}
	op.cpi.Broadcast(ocMsg)
}

// send a message to a specific replica
func (op *obcBatch) unicast(msgPayload []byte, receiverID uint64) (err error) {
	ocMsg := &pb.OpenchainMessage{
		Type:    pb.OpenchainMessage_CONSENSUS,
		Payload: msgPayload,
	}
	receiverHandle, err := getValidatorHandle(receiverID)
	if err != nil {
		return
	}
	return op.cpi.Unicast(ocMsg, receiverHandle)
}

func (op *obcBatch) sign(msg []byte) ([]byte, error) {
	return op.cpi.Sign(msg)
}

// verify message signature
func (op *obcBatch) verify(senderID uint64, signature []byte, message []byte) error {
	senderHandle, err := getValidatorHandle(senderID)
	if err != nil {
		return fmt.Errorf("Could not verify message from %v: %v", senderHandle.Name, err)
	}
	return op.cpi.Verify(senderHandle, signature, message)
}

// validate checks whether the request is valid syntactically.
// For now, we only need this for the obc-sieve verify/verify-set and flush messages.
// Thus, for obc-batch, this is a no-op.
func (op *obcBatch) validate(txRaw []byte) error {
	return nil
}

// execute an opaque request which corresponds to an OBC Transaction
func (op *obcBatch) execute(tbRaw []byte, rawMetadata []byte) {
	tb := &pb.TransactionBlock{}
	err := proto.Unmarshal(tbRaw, tb)
	if err != nil {
		return
	}

	txs := tb.Transactions
	txBatchID := base64.StdEncoding.EncodeToString(util.ComputeCryptoHash(tbRaw))

	for i, tx := range txs {
		txRaw, _ := proto.Marshal(tx)
		if err = op.validate(txRaw); err != nil {
			logger.Error("Request in transaction %d from batch %s did not verify: %s", i, txBatchID, err)
			return
		}
	}

	if err := op.cpi.BeginTxBatch(txBatchID); err != nil {
		logger.Error("Failed to begin transaction batch %s: %v", txBatchID, err)
		return
	}

	_, errs := op.cpi.ExecTXs(txs)
	if errs[len(txs)] != nil {
		logger.Error("Fail to execute transaction batch %s: %v", txBatchID, errs)
		if err = op.cpi.RollbackTxBatch(txBatchID); err != nil {
			panic(fmt.Errorf("Unable to rollback transaction batch %s: %v", txBatchID, err))
		}
		return
	}

	if err = op.cpi.CommitTxBatch(txBatchID, txs, nil, rawMetadata); err != nil {
		logger.Error("Failed to commit transaction batch %s to the ledger: %v", txBatchID, err)
		if err = op.cpi.RollbackTxBatch(txBatchID); err != nil {
			panic(fmt.Errorf("Unable to rollback transaction batch %s: %v", txBatchID, err))
		}
		return
	}
}

// signal when a view-change happened
func (op *obcBatch) viewChange(curView uint64) {
	if op.batchTimerActive {
		op.stopBatchTimer()
	}
}

// =============================================================================
// functions specific to batch mode
// =============================================================================

func (op *obcBatch) leaderProcReq(req *Request) error {
	digest := hashReq(req)
	op.batchStore[digest] = req

	if !op.batchTimerActive {
		op.startBatchTimer()
	}

	if len(op.batchStore) >= op.batchSize {
		op.sendBatch()
	}

	return nil
}

func (op *obcBatch) sendBatch() error {
	op.stopBatchTimer()
	// assemble new Request message
	txs := make([]*pb.Transaction, len(op.batchStore))
	var i int
	for d, req := range op.batchStore {
		txs[i] = &pb.Transaction{}
		err := proto.Unmarshal(req.Payload, txs[i])
		if err != nil {
			err = fmt.Errorf("Unable to unpack payload of request %d", i)
			logger.Error("%s", err)
			return err
		}
		i++
		delete(op.batchStore, d) // clean up
	}
	tb := &pb.TransactionBlock{Transactions: txs}
	tbPacked, err := proto.Marshal(tb)
	if err != nil {
		err = fmt.Errorf("Unable to pack transaction block for new batch request")
		logger.Error("%s", err)
		return err
	}
	// process internally
	op.pbft.request(tbPacked, op.pbft.id)
	// broadcast
	batchReq := &Request{Payload: tbPacked, ReplicaId: op.pbft.id}
	msg := &Message{&Message_Request{batchReq}}
	msgRaw, _ := proto.Marshal(msg)
	op.broadcast(msgRaw)

	return nil
}

// allow the primary to send a batch when the timer expires
func (op *obcBatch) batchTimerHander() {
	for {
		select {
		case <-op.batchTimer.C:
			op.pbft.lock.Lock()
			if op.pbft.closed {
				op.pbft.lock.Unlock()
				return
			}
			logger.Info("Replica %d batch timeout expired", op.pbft.id)
			if op.pbft.activeView && (len(op.batchStore) > 0) {
				op.sendBatch()
			}
			op.pbft.lock.Unlock()
		}
	}
}

func (op *obcBatch) startBatchTimer() {
	op.batchTimer.Reset(op.batchTimeout)
	logger.Debug("Replica %d started the batch timer", op.pbft.id)
	op.batchTimerActive = true
}

func (op *obcBatch) stopBatchTimer() {
	op.batchTimer.Stop()
	logger.Debug("Replica %d stopped the batch timer", op.pbft.id)
	op.batchTimerActive = false
	if op.pbft.closed {
		return
	}
loopBatch:
	for {
		select {
		case <-op.batchTimer.C:
		default:
			break loopBatch
		}
	}
}
<|endoftext|>


<commit_before>package network

import (
	"net"
	"testing"
	"time"
)

var chan1 = make(chan struct{})
var chan2 = make(chan struct{})

type Foo struct{}

func (f Foo) Bar(int32) error { chan1 <- struct{}{}; return nil }

func TestRegister(t *testing.T) {
	// create server
	tcps, err := NewTCPServer(9987)
	if err != nil {
		t.Fatal(err)
	}
	addr := NetAddress{"localhost", 9987}

	// register some handlers
	tcps.Register("Foo", func(net.Conn, []byte) error { chan2 <- struct{}{}; return nil })
	tcps.Register("Bar", new(Foo).Bar)

	// call them
	err = addr.RPC("Foo", 0, nil)
	if err != nil {
		t.Fatal(err)
	}
	err = addr.RPC("Bar", 0, nil)
	if err != nil {
		t.Fatal(err)
	}

	done := make(chan struct{})
	go func() {
		<-chan1
		<-chan2
		done <- struct{}{}
	}()

	// wait for messages to propagate
	select {
	// success
	case <-done:
		return

	// timeout
	case <-time.After(100 * time.Millisecond):
		t.Fatal("One or both handlers not called")
	}
}

func TestPeerSharing(t *testing.T) {
	// create server
	tcps, err := NewTCPServer(9981)
	if err != nil {
		t.Fatal(err)
	}

	// add a peer
	peer := NetAddress{"foo", 9001}
	tcps.AddPeer(peer)
	// tcps only has one peer, so RandomPeer() should return peer
	if tcps.RandomPeer() != peer {
		t.Fatal("server has bad peer list:", tcps.AddressBook())
	}

	// ask tcps for peers
	var resp []NetAddress
	err = tcps.myAddr.RPC("SharePeers", nil, &resp)
	if err != nil {
		t.Fatal(err)
	}
	// resp should be exactly []NetAddress{peer}
	if len(resp) != 1 || resp[0] != peer {
		t.Fatal("server gave bad peer list:", resp)
	}

	// add a couple more peers
	tcps.AddPeer(NetAddress{"bar", 9002})
	tcps.AddPeer(NetAddress{"baz", 9003})
	tcps.AddPeer(NetAddress{"quux", 9004})
	err = tcps.myAddr.RPC("SharePeers", nil, &resp)
	if err != nil {
		t.Fatal(err)
	}
	// resp should now contain 4 distinct addresses
	for i := 0; i < len(resp); i++ {
		for j := i + 1; j < len(resp); j++ {
			if resp[i] == resp[j] {
				t.Fatal("resp contains duplicate addresses:", resp)
			}
		}
	}
}
<commit_msg>add peer culling test<commit_after>package network

import (
	"net"
	"testing"
	"time"
)

var chan1 = make(chan struct{})
var chan2 = make(chan struct{})

type Foo struct{}

func (f Foo) Bar(int32) error { chan1 <- struct{}{}; return nil }

func TestRegister(t *testing.T) {
	// create server
	tcps, err := NewTCPServer(9987)
	if err != nil {
		t.Fatal(err)
	}
	addr := NetAddress{"localhost", 9987}

	// register some handlers
	tcps.Register("Foo", func(net.Conn, []byte) error { chan2 <- struct{}{}; return nil })
	tcps.Register("Bar", new(Foo).Bar)

	// call them
	err = addr.RPC("Foo", 0, nil)
	if err != nil {
		t.Fatal(err)
	}
	err = addr.RPC("Bar", 0, nil)
	if err != nil {
		t.Fatal(err)
	}

	done := make(chan struct{})
	go func() {
		<-chan1
		<-chan2
		done <- struct{}{}
	}()

	// wait for messages to propagate
	select {
	// success
	case <-done:
		return

	// timeout
	case <-time.After(100 * time.Millisecond):
		t.Fatal("One or both handlers not called")
	}
}

func TestPeerSharing(t *testing.T) {
	// create server
	tcps, err := NewTCPServer(9981)
	if err != nil {
		t.Fatal(err)
	}

	// add a peer
	peer := NetAddress{"foo", 9001}
	tcps.AddPeer(peer)
	// tcps only has one peer, so RandomPeer() should return peer
	if tcps.RandomPeer() != peer {
		t.Fatal("server has bad peer list:", tcps.AddressBook())
	}

	// ask tcps for peers
	var resp []NetAddress
	err = tcps.myAddr.RPC("SharePeers", nil, &resp)
	if err != nil {
		t.Fatal(err)
	}
	// resp should be exactly []NetAddress{peer}
	if len(resp) != 1 || resp[0] != peer {
		t.Fatal("server gave bad peer list:", resp)
	}

	// add a couple more peers
	tcps.AddPeer(NetAddress{"bar", 9002})
	tcps.AddPeer(NetAddress{"baz", 9003})
	tcps.AddPeer(NetAddress{"quux", 9004})
	err = tcps.myAddr.RPC("SharePeers", nil, &resp)
	if err != nil {
		t.Fatal(err)
	}
	// resp should now contain 4 distinct addresses
	for i := 0; i < len(resp); i++ {
		for j := i + 1; j < len(resp); j++ {
			if resp[i] == resp[j] {
				t.Fatal("resp contains duplicate addresses:", resp)
			}
		}
	}
}

func TestPeerCulling(t *testing.T) {
	// this test necessitates a timeout
	if testing.Short() {
		t.Skip()
	}

	// create server
	tcps, err := NewTCPServer(9005)
	if err != nil {
		t.Fatal(err)
	}

	// add google as a peer
	peer := NetAddress{"8.8.8.8", 9001}
	tcps.AddPeer(peer)

	// send a broadcast
	// doesn't need to be a real RPC
	tcps.Broadcast("QuestionWhoseAnswerIs", 42, nil)

	// peer should have been removed
	if len(tcps.AddressBook()) != 0 {
		t.Fatal("server did not remove dead peer:", tcps.AddressBook())
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
	"k8s.io/cloud-provider-openstack/pkg/kms/server"
)

var (
	socketpath  string
	cloudconfig string
)

func init() {
	flag.Set("logtostderr", "true")
}

func main() {
	flag.CommandLine.Parse([]string{})

	cmd := &cobra.Command{
		Use:   "barbican-kms-plugin",
		Short: "Barbican KMS plugin for kubernetes",
		RunE: func(cmd *cobra.Command, args []string) error {
			sigchan := make(chan os.Signal, 1)
			signal.Notify(sigchan, unix.SIGTERM, unix.SIGINT)
			err := server.Run(cloudconfig, socketpath, sigchan)
			return err
		},
	}

	cmd.Flags().AddGoFlagSet(flag.CommandLine)

	cmd.PersistentFlags().StringVar(&socketpath, "socketpath", "", "Barbican KMS Plugin unix socket endpoint")
	cmd.MarkPersistentFlagRequired("socketpath")

	cmd.PersistentFlags().StringVar(&cloudconfig, "cloud-config", "", "Barbican KMS Plugin cloud config")
	cmd.MarkPersistentFlagRequired("cloud-config")

	if err := cmd.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "%s", err.Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<commit_msg>Allow glog and klog to coexist for barbican<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"flag"
	"fmt"
	"os"
	"os/signal"

	"github.com/spf13/cobra"
	"golang.org/x/sys/unix"
	"k8s.io/apiserver/pkg/util/logs"
	"k8s.io/cloud-provider-openstack/pkg/kms/server"
	"k8s.io/klog"
)

var (
	socketpath  string
	cloudconfig string
)

func init() {
	flag.Set("logtostderr", "true")
}

func main() {
	// Glog requires this otherwise it complains.
	flag.CommandLine.Parse(nil)
	// This is a temporary hack to enable proper logging until upstream dependencies
	// are migrated to fully utilize klog instead of glog.
	klogFlags := flag.NewFlagSet("klog", flag.ExitOnError)
	klog.InitFlags(klogFlags)

	// Sync the glog and klog flags.
	flag.CommandLine.VisitAll(func(f1 *flag.Flag) {
		f2 := klogFlags.Lookup(f1.Name)
		if f2 != nil {
			value := f1.Value.String()
			f2.Value.Set(value)
		}
	})

	logs.InitLogs()
	defer logs.FlushLogs()

	cmd := &cobra.Command{
		Use:   "barbican-kms-plugin",
		Short: "Barbican KMS plugin for kubernetes",
		RunE: func(cmd *cobra.Command, args []string) error {
			sigchan := make(chan os.Signal, 1)
			signal.Notify(sigchan, unix.SIGTERM, unix.SIGINT)
			err := server.Run(cloudconfig, socketpath, sigchan)
			return err
		},
	}

	cmd.Flags().AddGoFlagSet(flag.CommandLine)

	cmd.PersistentFlags().StringVar(&socketpath, "socketpath", "", "Barbican KMS Plugin unix socket endpoint")
	cmd.MarkPersistentFlagRequired("socketpath")

	cmd.PersistentFlags().StringVar(&cloudconfig, "cloud-config", "", "Barbican KMS Plugin cloud config")
	cmd.MarkPersistentFlagRequired("cloud-config")

	if err := cmd.Execute(); err != nil {
		fmt.Fprintf(os.Stderr, "%s", err.Error())
		os.Exit(1)
	}

	os.Exit(0)
}
<|endoftext|>


<commit_before>package tests

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

	"github.com/sirupsen/logrus"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/portworx/torpedo/drivers/node"
	"github.com/portworx/torpedo/drivers/scheduler"
	. "github.com/portworx/torpedo/tests"

	// https://github.com/kubernetes/client-go/issues/242
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
)

const (
	scaleTimeout          = 10 * time.Minute
	nodeDeleteTimeoutMins = 7 * time.Minute
)

func TestASG(t *testing.T) {
	RegisterFailHandler(Fail)

	var specReporters []Reporter
	junitReporter := reporters.NewJUnitReporter("/testresults/junit_basic.xml")
	specReporters = append(specReporters, junitReporter)
	RunSpecsWithDefaultAndCustomReporters(t, "Torpedo : ASG", specReporters)
}

var _ = BeforeSuite(func() {
	InitInstance()
})

// This test performs basic test of scaling up and down the asg cluster
var _ = Describe("{ClusterScaleUpDown}", func() {
	It("has to validate that storage nodes are not lost during asg scaledown", func() {

		var contexts []*scheduler.Context
		for i := 0; i < Inst().ScaleFactor; i++ {
			contexts = append(contexts, ScheduleAndValidate(fmt.Sprintf("asgscaleupdown-%d", i))...)
		}

		intitialNodeCount, err := Inst().N.GetASGClusterSize()
		Expect(err).NotTo(HaveOccurred())

		scaleupCount := intitialNodeCount + intitialNodeCount/2
		Step(fmt.Sprintf("scale up cluster from %d to %d nodes and validate",
			intitialNodeCount, (scaleupCount/3)*3), func() {

			Scale(scaleupCount)
			Step(fmt.Sprintf("validate number of storage nodes after scale up"), func() {
				ValidateClusterSize(scaleupCount)
			})

		})

		Step(fmt.Sprintf("scale down cluster back to original size of %d nodes",
			intitialNodeCount), func() {
			Scale(intitialNodeCount)

			Step(fmt.Sprintf("wait for %s minutes for auto recovery of storeage nodes",
				Inst().AutoStorageNodeRecoveryTimeout.String()), func() {
				time.Sleep(Inst().AutoStorageNodeRecoveryTimeout)
			})

			// After scale down, get fresh list of nodes
			// by re-initializing scheduler and volume driver
			err = Inst().S.RefreshNodeRegistry()
			Expect(err).NotTo(HaveOccurred())

			err = Inst().V.RefreshDriverEndpoints()
			Expect(err).NotTo(HaveOccurred())

			Step(fmt.Sprintf("validate number of storage nodes after scale down"), func() {
				ValidateClusterSize(intitialNodeCount)
			})
		})

		opts := make(map[string]bool)
		opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
		ValidateAndDestroy(contexts, opts)
	})
})

// This test randomly kills one volume driver node and ensures cluster remains
// intact by ASG
var _ = Describe("{ASGKillRandomNodes}", func() {
	It("keeps killing worker nodes", func() {

		var contexts []*scheduler.Context
		var err error

		// Get list of nodes where storage driver is installed
		storageDriverNodes := node.GetStorageDriverNodes()
		Expect(err).NotTo(HaveOccurred())

		Step("Ensure apps are deployed", func() {
			for i := 0; i < Inst().ScaleFactor; i++ {
				contexts = append(contexts, ScheduleAndValidate(fmt.Sprintf("asgchaos-%d", i))...)
			}
		})

		Step("Randomly kill one storage node", func() {

			// set frequency mins depending on the chaos level
			var frequency int
			switch Inst().ChaosLevel {
			case 5:
				frequency = 15
			case 4:
				frequency = 30
			case 3:
				frequency = 45
			case 2:
				frequency = 60
			case 1:
				frequency = 90
			default:
				frequency = 30

			}
			if Inst().MinRunTimeMins == 0 {
				// Run once
				KillANodeAndValidate(storageDriverNodes)

				// Validate applications and tear down
				opts := make(map[string]bool)
				opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
				ValidateAndDestroy(contexts, opts)
			} else {
				// Run once till timer gets triggered
				KillANodeAndValidate(storageDriverNodes)

				Step("validate applications", func() {
					for _, ctx := range contexts {
						ValidateContext(ctx)
					}
				})

				// Run repeatedly
				ticker := time.NewTicker(time.Duration(frequency) * time.Minute)
				stopChannel := time.After(time.Duration(Inst().MinRunTimeMins) * time.Minute)
			L:
				for {
					select {
					case <-ticker.C:
						KillANodeAndValidate(storageDriverNodes)

						Step("validate applications", func() {
							for _, ctx := range contexts {
								ValidateContext(ctx)
							}
						})
					case <-stopChannel:
						ticker.Stop()
						// ticker may expire/time out in between, apps may not be
						// in correct condition to be validated. Just tear them down.
						opts := make(map[string]bool)
						opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
						Step("destroy apps", func() {
							for _, ctx := range contexts {
								TearDownContext(ctx, opts)
							}
						})
						break L
					}
				}
			}
		})
	})
})

var _ = AfterSuite(func() {
	PerformSystemCheck()
	ValidateCleanup()
})

func init() {
	ParseFlags()
}

func Scale(count int64) {
	// In multi-zone ASG cluster, node count is per zone
	perZoneCount := count / 3

	err := Inst().N.SetASGClusterSize(perZoneCount, scaleTimeout)
	Expect(err).NotTo(HaveOccurred())
}

func ValidateClusterSize(count int64) {
	// In multi-zone ASG cluster, node count is per zone
	perZoneCount := count / 3

	// Validate total node count
	currentNodeCount, err := Inst().N.GetASGClusterSize()
	Expect(err).NotTo(HaveOccurred())
	Expect(currentNodeCount).Should(Equal(perZoneCount * 3))

	// Validate storage node count
	var expectedStorageNodesPerZone int
	if Inst().MaxStorageNodesPerAZ <= int(perZoneCount) {
		expectedStorageNodesPerZone = Inst().MaxStorageNodesPerAZ
	} else {
		expectedStorageNodesPerZone = int(perZoneCount)
	}
	storageNodes, err := getStorageNodes()
	Expect(err).NotTo(HaveOccurred())
	Expect(len(storageNodes)).Should(Equal(expectedStorageNodesPerZone * 3))
	logrus.Infof("Validated successfully that [%d] storage nodes are present", len(storageNodes))
}

func getStorageNodes() ([]node.Node, error) {

	storageNodes := []node.Node{}
	nodes := node.GetStorageDriverNodes()

	for _, node := range nodes {
		devices, err := Inst().V.GetStorageDevices(node)
		if err != nil {
			return nil, err
		}
		if len(devices) > 0 {
			storageNodes = append(storageNodes, node)
		}
	}
	return storageNodes, nil
}

func KillANodeAndValidate(storageDriverNodes []node.Node) {
	rand.Seed(time.Now().Unix())
	nodeToKill := storageDriverNodes[rand.Intn(len(storageDriverNodes))]

	Step(fmt.Sprintf("Deleting node [%v]", nodeToKill.Name), func() {
		err := Inst().N.DeleteNode(nodeToKill, nodeDeleteTimeoutMins)
		Expect(err).NotTo(HaveOccurred())
	})

	Step("Wait for 10 min. to node get replaced by autoscalling group", func() {
		time.Sleep(10 * time.Minute)
	})

	err := Inst().S.RefreshNodeRegistry()
	Expect(err).NotTo(HaveOccurred())

	err = Inst().V.RefreshDriverEndpoints()
	Expect(err).NotTo(HaveOccurred())

	Step(fmt.Sprintf("Validate number of storage nodes after killing node [%v]", nodeToKill.Name), func() {
		ValidateClusterSize(int64(len(storageDriverNodes)))
	})
}
<commit_msg>* Add clear print statements in case of gomega assertion failures (#364)<commit_after>package tests

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

	"github.com/sirupsen/logrus"

	. "github.com/onsi/ginkgo"
	"github.com/onsi/ginkgo/reporters"
	. "github.com/onsi/gomega"
	"github.com/portworx/torpedo/drivers/node"
	"github.com/portworx/torpedo/drivers/scheduler"
	. "github.com/portworx/torpedo/tests"

	// https://github.com/kubernetes/client-go/issues/242
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"
)

const (
	scaleTimeout          = 10 * time.Minute
	nodeDeleteTimeoutMins = 7 * time.Minute
)

func TestASG(t *testing.T) {
	RegisterFailHandler(Fail)

	var specReporters []Reporter
	junitReporter := reporters.NewJUnitReporter("/testresults/junit_basic.xml")
	specReporters = append(specReporters, junitReporter)
	RunSpecsWithDefaultAndCustomReporters(t, "Torpedo : ASG", specReporters)
}

var _ = BeforeSuite(func() {
	InitInstance()
})

// This test performs basic test of scaling up and down the asg cluster
var _ = Describe("{ClusterScaleUpDown}", func() {
	It("has to validate that storage nodes are not lost during asg scaledown", func() {

		var contexts []*scheduler.Context
		for i := 0; i < Inst().ScaleFactor; i++ {
			contexts = append(contexts, ScheduleAndValidate(fmt.Sprintf("asgscaleupdown-%d", i))...)
		}

		intitialNodeCount, err := Inst().N.GetASGClusterSize()
		Expect(err).NotTo(HaveOccurred())

		scaleupCount := intitialNodeCount + intitialNodeCount/2
		Step(fmt.Sprintf("scale up cluster from %d to %d nodes and validate",
			intitialNodeCount, (scaleupCount/3)*3), func() {

			Scale(scaleupCount)
			Step(fmt.Sprintf("validate number of storage nodes after scale up"), func() {
				ValidateClusterSize(scaleupCount)
			})

		})

		Step(fmt.Sprintf("scale down cluster back to original size of %d nodes",
			intitialNodeCount), func() {
			Scale(intitialNodeCount)

			Step(fmt.Sprintf("wait for %s minutes for auto recovery of storeage nodes",
				Inst().AutoStorageNodeRecoveryTimeout.String()), func() {
				time.Sleep(Inst().AutoStorageNodeRecoveryTimeout)
			})

			// After scale down, get fresh list of nodes
			// by re-initializing scheduler and volume driver
			err = Inst().S.RefreshNodeRegistry()
			Expect(err).NotTo(HaveOccurred())

			err = Inst().V.RefreshDriverEndpoints()
			Expect(err).NotTo(HaveOccurred())

			Step(fmt.Sprintf("validate number of storage nodes after scale down"), func() {
				ValidateClusterSize(intitialNodeCount)
			})
		})

		opts := make(map[string]bool)
		opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
		ValidateAndDestroy(contexts, opts)
	})
})

// This test randomly kills one volume driver node and ensures cluster remains
// intact by ASG
var _ = Describe("{ASGKillRandomNodes}", func() {
	It("keeps killing worker nodes", func() {

		var contexts []*scheduler.Context
		var err error

		// Get list of nodes where storage driver is installed
		storageDriverNodes := node.GetStorageDriverNodes()
		Expect(err).NotTo(HaveOccurred())

		Step("Ensure apps are deployed", func() {
			for i := 0; i < Inst().ScaleFactor; i++ {
				contexts = append(contexts, ScheduleAndValidate(fmt.Sprintf("asgchaos-%d", i))...)
			}
		})

		Step("Randomly kill one storage node", func() {

			// set frequency mins depending on the chaos level
			var frequency int
			switch Inst().ChaosLevel {
			case 5:
				frequency = 15
			case 4:
				frequency = 30
			case 3:
				frequency = 45
			case 2:
				frequency = 60
			case 1:
				frequency = 90
			default:
				frequency = 30

			}
			if Inst().MinRunTimeMins == 0 {
				// Run once
				KillANodeAndValidate(storageDriverNodes)

				// Validate applications and tear down
				opts := make(map[string]bool)
				opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
				ValidateAndDestroy(contexts, opts)
			} else {
				// Run once till timer gets triggered
				KillANodeAndValidate(storageDriverNodes)

				Step("validate applications", func() {
					for _, ctx := range contexts {
						ValidateContext(ctx)
					}
				})

				// Run repeatedly
				ticker := time.NewTicker(time.Duration(frequency) * time.Minute)
				stopChannel := time.After(time.Duration(Inst().MinRunTimeMins) * time.Minute)
			L:
				for {
					select {
					case <-ticker.C:
						KillANodeAndValidate(storageDriverNodes)

						Step("validate applications", func() {
							for _, ctx := range contexts {
								ValidateContext(ctx)
							}
						})
					case <-stopChannel:
						ticker.Stop()
						// ticker may expire/time out in between, apps may not be
						// in correct condition to be validated. Just tear them down.
						opts := make(map[string]bool)
						opts[scheduler.OptionsWaitForResourceLeakCleanup] = true
						Step("destroy apps", func() {
							for _, ctx := range contexts {
								TearDownContext(ctx, opts)
							}
						})
						break L
					}
				}
			}
		})
	})
})

var _ = AfterSuite(func() {
	PerformSystemCheck()
	ValidateCleanup()
})

func init() {
	ParseFlags()
}

func Scale(count int64) {
	// In multi-zone ASG cluster, node count is per zone
	perZoneCount := count / 3

	err := Inst().N.SetASGClusterSize(perZoneCount, scaleTimeout)
	Expect(err).NotTo(HaveOccurred())
}

func ValidateClusterSize(count int64) {
	// In multi-zone ASG cluster, node count is per zone
	perZoneCount := count / 3

	// Validate total node count
	currentNodeCount, err := Inst().N.GetASGClusterSize()
	Expect(err).NotTo(HaveOccurred())
	Expect(perZoneCount*3).Should(Equal(currentNodeCount),
		"ASG cluster size is not as expected."+
			" Current size is [%d]. Expected ASG size is [%d]",
		currentNodeCount, perZoneCount*3)

	// Validate storage node count
	var expectedStorageNodesPerZone int
	if Inst().MaxStorageNodesPerAZ <= int(perZoneCount) {
		expectedStorageNodesPerZone = Inst().MaxStorageNodesPerAZ
	} else {
		expectedStorageNodesPerZone = int(perZoneCount)
	}
	storageNodes, err := getStorageNodes()
	Expect(err).NotTo(HaveOccurred())
	Expect(len(storageNodes)).Should(Equal(expectedStorageNodesPerZone*3),
		"Current number of storeage nodes [%d] does not match expected number of storage nodes [%d]",
		len(storageNodes), expectedStorageNodesPerZone*3)

	logrus.Infof("Validated successfully that [%d] storage nodes are present", len(storageNodes))
}

func getStorageNodes() ([]node.Node, error) {

	storageNodes := []node.Node{}
	nodes := node.GetStorageDriverNodes()

	for _, node := range nodes {
		devices, err := Inst().V.GetStorageDevices(node)
		if err != nil {
			return nil, err
		}
		if len(devices) > 0 {
			storageNodes = append(storageNodes, node)
		}
	}
	return storageNodes, nil
}

func KillANodeAndValidate(storageDriverNodes []node.Node) {
	rand.Seed(time.Now().Unix())
	nodeToKill := storageDriverNodes[rand.Intn(len(storageDriverNodes))]

	Step(fmt.Sprintf("Deleting node [%v]", nodeToKill.Name), func() {
		err := Inst().N.DeleteNode(nodeToKill, nodeDeleteTimeoutMins)
		Expect(err).NotTo(HaveOccurred())
	})

	Step("Wait for 10 min. to node get replaced by autoscalling group", func() {
		time.Sleep(10 * time.Minute)
	})

	err := Inst().S.RefreshNodeRegistry()
	Expect(err).NotTo(HaveOccurred())

	err = Inst().V.RefreshDriverEndpoints()
	Expect(err).NotTo(HaveOccurred())

	Step(fmt.Sprintf("Validate number of storage nodes after killing node [%v]", nodeToKill.Name), func() {
		ValidateClusterSize(int64(len(storageDriverNodes)))
	})
}
<|endoftext|>


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

package checkers_test

import (
	"testing"

	. "launchpad.net/gocheck"
	. "launchpad.net/juju-core/testing/checkers"
)

func Test(t *testing.T) { TestingT(t) }

type CheckerSuite struct{}

var _ = Suite(&CheckerSuite{})

func (s *CheckerSuite) TestHasPrefix(c *C) {
	c.Assert("foo bar", HasPrefix, "foo")
	c.Assert("foo bar", Not(HasPrefix), "omg")
}

func (s *CheckerSuite) TestHasSuffix(c *C) {
	c.Assert("foo bar", HasSuffix, "bar")
	c.Assert("foo bar", Not(HasSuffix), "omg")
}

func (s *CheckerSuite) TestContains(c *C) {
	c.Assert("foo bar baz", Contains, "foo")
	c.Assert("foo bar baz", Contains, "bar")
	c.Assert("foo bar baz", Contains, "baz")
	c.Assert("foo bar baz", Not(Contains), "omg")
}
<commit_msg>Update for new style<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package checkers_test

import (
	"testing"

	gc "launchpad.net/gocheck"
	jc "launchpad.net/juju-core/testing/checkers"
)

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

type CheckerSuite struct{}

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

func (s *CheckerSuite) TestHasPrefix(c *gc.C) {
	c.Assert("foo bar", jc.HasPrefix, "foo")
	c.Assert("foo bar", gc.Not(jc.HasPrefix), "omg")
}

func (s *CheckerSuite) TestHasSuffix(c *gc.C) {
	c.Assert("foo bar", jc.HasSuffix, "bar")
	c.Assert("foo bar", gc.Not(jc.HasSuffix), "omg")
}

func (s *CheckerSuite) TestContains(c *gc.C) {
	c.Assert("foo bar baz", jc.Contains, "foo")
	c.Assert("foo bar baz", jc.Contains, "bar")
	c.Assert("foo bar baz", jc.Contains, "baz")
	c.Assert("foo bar baz", gc.Not(jc.Contains), "omg")
}
<|endoftext|>


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

package state_test

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/errors"
	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/network"
	"github.com/juju/juju/state"
)

type IPAddressSuite struct {
	ConnSuite
}

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

func (s *IPAddressSuite) SetUpTest(c *gc.C) {
	s.ConnSuite.SetUpTest(c)
	s.policy.GetConstraintsValidator = func(*config.Config) (constraints.Validator, error) {
		validator := constraints.NewValidator()
		validator.RegisterConflicts([]string{constraints.InstanceType}, []string{constraints.Mem})
		validator.RegisterUnsupported([]string{constraints.CpuPower})
		return validator, nil
	}
}

func (s *IPAddressSuite) TestAddIPAddress(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	assertAddress := func(ipAddr *state.IPAddress) {
		c.Assert(ipAddr.Value(), gc.Equals, "192.168.1.0")
		c.Assert(ipAddr.SubnetId(), gc.Equals, "foobar")
		c.Assert(ipAddr.Type(), gc.Equals, addr.Type)
		c.Assert(ipAddr.Scope(), gc.Equals, network.ScopePublic)
		c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)
	}
	checkAddress(ipAddr)

	// verify the address was stored in the state
	ipAddr, err = s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	checkAddress(ipAddr)
}

func (s *IPAddressSuite) TestAddIPAddressInvalid(c *gc.C) {
	addr := network.Address{Value: "foo"}
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(errors.Cause(err), gc.ErrorMatches, `invalid IP address "foo"`)
}

func (s *IPAddressSuite) TestAddIPAddressAlreadyExists(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	_, err = s.State.AddIPAddress(addr, "foobar")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsAlreadyExists)
}

func (s *IPAddressSuite) TestIPAddressNotFound(c *gc.C) {
	_, err := s.State.IPAddress("192.168.1.0")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsNotFound)
}

func (s *IPAddressSuite) TestIPAddressRemove(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	_, err = s.State.IPAddress("192.168.1.0")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsNotFound)
}

func (s *IPAddressSuite) TestIPAddressSetState(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)

	err = ipAddr.SetState(state.AddressStateAllocated)
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.State(), gc.Equals, state.AddressStateAllocated)

	// setting the state to the same state is permitted
	err = ipAddr.SetState(state.AddressStateAllocated)
	c.Assert(err, jc.ErrorIsNil)

	// setting back to unknown isn't permitted
	err = ipAddr.SetState(state.AddressStateUnknown)
	c.Assert(err, gc.ErrorMatches, `cannot set IP address .* to state "".*`)
}

func (s *IPAddressSuite) TestIPAddressSetMachineId(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.MachineId(), gc.Equals, "")

	err = ipAddr.SetMachineId("wibble")
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.MachineId(), gc.Equals, "wibble")
}

func (s *IPAddressSuite) TestIPAddressSetInterfaceId(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "")

	err = ipAddr.SetInterfaceId("wobble")
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.InterfaceId(), gc.Equals, "wobble")
}

func (s *IPAddressSuite) TestIPAddressAddress(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.Address(), gc.Equals, addr)

}
<commit_msg>Fix missed rename<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/errors"
	"github.com/juju/juju/constraints"
	"github.com/juju/juju/environs/config"
	"github.com/juju/juju/network"
	"github.com/juju/juju/state"
)

type IPAddressSuite struct {
	ConnSuite
}

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

func (s *IPAddressSuite) SetUpTest(c *gc.C) {
	s.ConnSuite.SetUpTest(c)
	s.policy.GetConstraintsValidator = func(*config.Config) (constraints.Validator, error) {
		validator := constraints.NewValidator()
		validator.RegisterConflicts([]string{constraints.InstanceType}, []string{constraints.Mem})
		validator.RegisterUnsupported([]string{constraints.CpuPower})
		return validator, nil
	}
}

func (s *IPAddressSuite) TestAddIPAddress(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	assertAddress := func(ipAddr *state.IPAddress) {
		c.Assert(ipAddr.Value(), gc.Equals, "192.168.1.0")
		c.Assert(ipAddr.SubnetId(), gc.Equals, "foobar")
		c.Assert(ipAddr.Type(), gc.Equals, addr.Type)
		c.Assert(ipAddr.Scope(), gc.Equals, network.ScopePublic)
		c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)
	}
	assertAddress(ipAddr)

	// verify the address was stored in the state
	ipAddr, err = s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	assertAddress(ipAddr)
}

func (s *IPAddressSuite) TestAddIPAddressInvalid(c *gc.C) {
	addr := network.Address{Value: "foo"}
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(errors.Cause(err), gc.ErrorMatches, `invalid IP address "foo"`)
}

func (s *IPAddressSuite) TestAddIPAddressAlreadyExists(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	_, err = s.State.AddIPAddress(addr, "foobar")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsAlreadyExists)
}

func (s *IPAddressSuite) TestIPAddressNotFound(c *gc.C) {
	_, err := s.State.IPAddress("192.168.1.0")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsNotFound)
}

func (s *IPAddressSuite) TestIPAddressRemove(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	_, err = s.State.IPAddress("192.168.1.0")
	c.Assert(errors.Cause(err), jc.Satisfies, errors.IsNotFound)
}

func (s *IPAddressSuite) TestIPAddressSetState(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)

	err = ipAddr.SetState(state.AddressStateAllocated)
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.State(), gc.Equals, state.AddressStateAllocated)

	// setting the state to the same state is permitted
	err = ipAddr.SetState(state.AddressStateAllocated)
	c.Assert(err, jc.ErrorIsNil)

	// setting back to unknown isn't permitted
	err = ipAddr.SetState(state.AddressStateUnknown)
	c.Assert(err, gc.ErrorMatches, `cannot set IP address .* to state "".*`)
}

func (s *IPAddressSuite) TestIPAddressSetMachineId(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.MachineId(), gc.Equals, "")

	err = ipAddr.SetMachineId("wibble")
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.MachineId(), gc.Equals, "wibble")
}

func (s *IPAddressSuite) TestIPAddressSetInterfaceId(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "")

	err = ipAddr.SetInterfaceId("wobble")
	c.Assert(err, jc.ErrorIsNil)

	freshCopy, err := s.State.IPAddress("192.168.1.0")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.InterfaceId(), gc.Equals, "wobble")
}

func (s *IPAddressSuite) TestIPAddressAddress(c *gc.C) {
	addr := network.NewAddress("192.168.1.0", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.Address(), gc.Equals, addr)

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var packages int // packages to update

func main() {
	sock, err := net.Listen("unix", "/tmp/packagesd.sock")
	check(err)
	defer func() { check(sock.Close()) }()
	check(os.Chmod("/tmp/packagesd.sock", 0666))

	getPackages()
	go func() {
		for {
			time.Sleep(time.Minute * 10)
			getPackages()
		}
	}()

	go func() {
		for {
			conn, err := sock.Accept()
			check(err)
			_, _ = conn.Write([]byte(strconv.Itoa(packages)))
			check(conn.Close())
		}
	}()

	// gracefully exit on interrupt, closing the socket
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)
	<-interrupt
}

func check(err error) {
	if err != nil {
		panic(err)
	}
}

func getPackages() {
	packages = 0

	_, err := exec.LookPath("pacaur")
	if err == nil {
		queryPacaur()
		return
	}

	_, err = exec.LookPath("apt")
	if err == nil {
		queryApt()
		return
	}
}

func queryPacaur() {
	// update pacman databases
	check(exec.Command("pacman", "-Syy").Run())

	// count packages with available updates
	cmd := exec.Command("pacaur", "-Qu")
	stdout, err := cmd.StdoutPipe()
	check(err)
	check(cmd.Start())

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		packages++
	}
	check(scanner.Err())
	_ = cmd.Wait()
}

func queryApt() {
	// update apt databases
	check(exec.Command("apt", "update").Run())

	// count packages with available updates
	cmd := exec.Command("apt", "list", "--upgradable")
	stdout, err := cmd.StdoutPipe()
	check(err)
	check(cmd.Start())

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		packages++
	}
	// account for the first line, "Listing... Done"
	packages--
	check(scanner.Err())
	_ = cmd.Wait()
}
<commit_msg>added support for dnf<commit_after>package main

import (
	"bufio"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var packages int // packages to update

func main() {
	sock, err := net.Listen("unix", "/tmp/packagesd.sock")
	check(err)
	defer func() { check(sock.Close()) }()
	check(os.Chmod("/tmp/packagesd.sock", 0666))

	getPackages()
	go func() {
		for {
			time.Sleep(time.Minute * 10)
			getPackages()
		}
	}()

	go func() {
		for {
			conn, err := sock.Accept()
			check(err)
			_, _ = conn.Write([]byte(strconv.Itoa(packages)))
			check(conn.Close())
		}
	}()

	// gracefully exit on interrupt, closing the socket
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, syscall.SIGINT, syscall.SIGTERM)
	<-interrupt
}

func check(err error) {
	if err != nil {
		panic(err)
	}
}

func getPackages() {
	packages = 0

	_, err := exec.LookPath("pacaur")
	if err == nil {
		queryPacaur()
		return
	}

	_, err = exec.LookPath("dnf")
	if err == nil {
		queryDnf()
		return
	}

	_, err = exec.LookPath("apt")
	if err == nil {
		queryApt()
		return
	}
}

func queryPacaur() {
	// update pacman databases
	check(exec.Command("pacman", "-Syy").Run())

	// count packages with available updates
	cmd := exec.Command("pacaur", "-Qu")
	stdout, err := cmd.StdoutPipe()
	check(err)
	check(cmd.Start())

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		packages++
	}
	check(scanner.Err())
	_ = cmd.Wait()
}

func queryDnf() {
	// count packages with available updates
	cmd := exec.Command("dnf", "check-update", "-q")
	stdout, err := cmd.StdoutPipe()
	check(err)
	check(cmd.Start())

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		packages++
	}
	// account for the blank line
	packages--
	check(scanner.Err())
	_ = cmd.Wait()
}

func queryApt() {
	// update apt databases
	check(exec.Command("apt", "update").Run())

	// count packages with available updates
	cmd := exec.Command("apt", "list", "--upgradable")
	stdout, err := cmd.StdoutPipe()
	check(err)
	check(cmd.Start())

	scanner := bufio.NewScanner(stdout)
	for scanner.Scan() {
		packages++
	}
	// account for the first line, "Listing... Done"
	packages--
	check(scanner.Err())
	_ = cmd.Wait()
}
<|endoftext|>


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

package state_test

import (
	"fmt"

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

	"github.com/juju/juju/network"
	"github.com/juju/juju/state"
)

type IPAddressSuite struct {
	ConnSuite
}

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

func (s *IPAddressSuite) assertAddress(
	c *gc.C,
	ipAddr *state.IPAddress,
	addr network.Address,
	ipState state.AddressState,
	machineId, ifaceId, subnetId string,
) {
	c.Assert(ipAddr, gc.NotNil)
	c.Assert(ipAddr.MachineId(), gc.Equals, machineId)
	c.Assert(ipAddr.InterfaceId(), gc.Equals, ifaceId)
	c.Assert(ipAddr.SubnetId(), gc.Equals, subnetId)
	c.Assert(ipAddr.Value(), gc.Equals, addr.Value)
	c.Assert(ipAddr.Type(), gc.Equals, addr.Type)
	c.Assert(ipAddr.Scope(), gc.Equals, addr.Scope)
	c.Assert(ipAddr.State(), gc.Equals, ipState)
	c.Assert(ipAddr.Address(), jc.DeepEquals, addr)
	c.Assert(ipAddr.String(), gc.Equals, addr.String())
}

func (s *IPAddressSuite) TestAddIPAddress(c *gc.C) {
	for i, test := range []string{"0.1.2.3", "2001:db8::1"} {
		c.Logf("test %d: %q", i, test)
		addr := network.NewAddress(test, network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		s.assertAddress(c, ipAddr, addr, state.AddressStateUnknown, "", "", "foobar")
		c.Assert(ipAddr.Life(), gc.Equals, state.Alive)

		// verify the address was stored in the state
		ipAddr, err = s.State.IPAddress(test)
		c.Assert(err, jc.ErrorIsNil)
		s.assertAddress(c, ipAddr, addr, state.AddressStateUnknown, "", "", "foobar")
	}
}

func (s *IPAddressSuite) TestAddIPAddressInvalid(c *gc.C) {
	addr := network.Address{Value: "foo"}
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.Satisfies, errors.IsNotValid)
	c.Assert(err, gc.ErrorMatches, `cannot add IP address "foo": address not valid`)
}

func (s *IPAddressSuite) TestAddIPAddressAlreadyExists(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	_, err = s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.Satisfies, errors.IsAlreadyExists)
	c.Assert(err, gc.ErrorMatches,
		`cannot add IP address "public:0.1.2.3": address already exists`,
	)
}

func (s *IPAddressSuite) TestIPAddressNotFound(c *gc.C) {
	_, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
	c.Assert(err, gc.ErrorMatches, `IP address "0.1.2.3" not found`)
}

func (s *IPAddressSuite) TestEnsureDeadRemove(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	// Should not be able to remove an Alive IP address.
	c.Assert(ipAddr.Life(), gc.Equals, state.Alive)
	err = ipAddr.Remove()
	msg := fmt.Sprintf("cannot remove IP address %q: IP address is not dead", ipAddr.String())
	c.Assert(err, gc.ErrorMatches, msg)

	err = ipAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	// EnsureDead twice should not be an error
	err = ipAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	// Remove twice is also fine.
	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	_, err = s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
	c.Assert(err, gc.ErrorMatches, `IP address "0.1.2.3" not found`)
}

func (s *IPAddressSuite) TestSetStateDead(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	copyIPAddr, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	err = copyIPAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.SetState(state.AddressStateAllocated)
	msg := fmt.Sprintf(`cannot set IP address %q to state "allocated": address is dead`, ipAddr.String())
	c.Assert(err, gc.ErrorMatches, msg)
}

func (s *IPAddressSuite) TestAllocateToDead(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	copyIPAddr, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	err = copyIPAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	msg := fmt.Sprintf(`cannot allocate IP address %q to machine "foobar", interface "wibble": address is dead`, ipAddr.String())
	err = ipAddr.AllocateTo("foobar", "wibble")
	c.Assert(err, gc.ErrorMatches, msg)
}

func (s *IPAddressSuite) TestAddressStateString(c *gc.C) {
	for i, test := range []struct {
		ipState state.AddressState
		expect  string
	}{{
		state.AddressStateUnknown,
		"<unknown>",
	}, {
		state.AddressStateAllocated,
		"allocated",
	}, {
		state.AddressStateUnavailable,
		"unavailable",
	}} {
		c.Logf("test %d: %q -> %q", i, test.ipState, test.expect)
		c.Check(test.ipState.String(), gc.Equals, test.expect)
	}
}

func (s *IPAddressSuite) TestSetState(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)

	for i, test := range []struct {
		initial, changeTo state.AddressState
		err               string
	}{{
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateUnknown,
	}, {
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateAllocated,
	}, {
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateUnavailable,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateAllocated,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateUnavailable,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateUnknown,
		err: `cannot set IP address "public:0.1.2.3" to state "<unknown>": ` +
			`transition from "allocated" not valid`,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateUnknown,
		err: `cannot set IP address "public:0.1.2.3" to state "<unknown>": ` +
			`transition from "unavailable" not valid`,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateUnavailable,
		err: `cannot set IP address "public:0.1.2.3" to state "unavailable": ` +
			`transition from "allocated" not valid`,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateAllocated,
		err: `cannot set IP address "public:0.1.2.3" to state "allocated": ` +
			`transition from "unavailable" not valid`,
	}} {
		c.Logf("test %d: %q -> %q ok:%v", i, test.initial, test.changeTo, test.err == "")
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Check(err, jc.ErrorIsNil)

		// Initially, all addresses have AddressStateUnknown.
		c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)

		if test.initial != state.AddressStateUnknown {
			err = ipAddr.SetState(test.initial)
			c.Check(err, jc.ErrorIsNil)
		}
		err = ipAddr.SetState(test.changeTo)
		if test.err != "" {
			c.Check(err, gc.ErrorMatches, test.err)
			c.Check(err, jc.Satisfies, errors.IsNotValid)
			c.Check(ipAddr.EnsureDead(), jc.ErrorIsNil)
			c.Check(ipAddr.Remove(), jc.ErrorIsNil)
			continue
		}
		c.Check(err, jc.ErrorIsNil)
		c.Check(ipAddr.State(), gc.Equals, test.changeTo)
		c.Check(ipAddr.EnsureDead(), jc.ErrorIsNil)
		c.Check(ipAddr.Remove(), jc.ErrorIsNil)
	}
}

func (s *IPAddressSuite) TestAllocateTo(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)
	c.Assert(ipAddr.MachineId(), gc.Equals, "")
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "")

	err = ipAddr.AllocateTo("wibble", "wobble")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateAllocated)
	c.Assert(ipAddr.MachineId(), gc.Equals, "wibble")
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "wobble")

	freshCopy, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.State(), gc.Equals, state.AddressStateAllocated)
	c.Assert(freshCopy.MachineId(), gc.Equals, "wibble")
	c.Assert(freshCopy.InterfaceId(), gc.Equals, "wobble")

	// allocating twice should fail.
	err = ipAddr.AllocateTo("m", "i")
	c.Assert(err, gc.ErrorMatches,
		`cannot allocate IP address "public:0.1.2.3" to machine "m", interface "i": `+
			`already allocated or unavailable`,
	)
}

func (s *IPAddressSuite) TestAddress(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.Address(), jc.DeepEquals, addr)

}

func (s *IPAddressSuite) TestAllocatedIPAddresses(c *gc.C) {
	addresses := [][]string{
		{"0.1.2.3", "wibble"},
		{"0.1.2.4", "wibble"},
		{"0.1.2.5", "wobble"},
	}
	for _, details := range addresses {
		addr := network.NewAddress(details[0], network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		err = ipAddr.AllocateTo(details[1], "wobble")
		c.Assert(err, jc.ErrorIsNil)
	}
	result, err := s.State.AllocatedIPAddresses("wibble")
	c.Assert(err, jc.ErrorIsNil)
	addr1, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	addr2, err := s.State.IPAddress("0.1.2.4")
	c.Assert(err, jc.ErrorIsNil)
	expected := []*state.IPAddress{addr1, addr2}
	c.Assert(result, jc.SameContents, expected)

}

func (s *IPAddressSuite) TestRefresh(c *gc.C) {
	rawAddr := network.NewAddress("0.1.2.3", network.ScopeUnknown)
	addr, err := s.State.AddIPAddress(rawAddr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	addrCopy, err := s.State.IPAddress(rawAddr.Value)
	c.Assert(err, jc.ErrorIsNil)

	err = addr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	c.Assert(addrCopy.Life(), gc.Equals, state.Alive)
	err = addrCopy.Refresh()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(addrCopy.Life(), gc.Equals, state.Dead)
}
<commit_msg>Test DeadIPAddresses<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package state_test

import (
	"fmt"

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

	"github.com/juju/juju/network"
	"github.com/juju/juju/state"
)

type IPAddressSuite struct {
	ConnSuite
}

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

func (s *IPAddressSuite) assertAddress(
	c *gc.C,
	ipAddr *state.IPAddress,
	addr network.Address,
	ipState state.AddressState,
	machineId, ifaceId, subnetId string,
) {
	c.Assert(ipAddr, gc.NotNil)
	c.Assert(ipAddr.MachineId(), gc.Equals, machineId)
	c.Assert(ipAddr.InterfaceId(), gc.Equals, ifaceId)
	c.Assert(ipAddr.SubnetId(), gc.Equals, subnetId)
	c.Assert(ipAddr.Value(), gc.Equals, addr.Value)
	c.Assert(ipAddr.Type(), gc.Equals, addr.Type)
	c.Assert(ipAddr.Scope(), gc.Equals, addr.Scope)
	c.Assert(ipAddr.State(), gc.Equals, ipState)
	c.Assert(ipAddr.Address(), jc.DeepEquals, addr)
	c.Assert(ipAddr.String(), gc.Equals, addr.String())
}

func (s *IPAddressSuite) TestAddIPAddress(c *gc.C) {
	for i, test := range []string{"0.1.2.3", "2001:db8::1"} {
		c.Logf("test %d: %q", i, test)
		addr := network.NewAddress(test, network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		s.assertAddress(c, ipAddr, addr, state.AddressStateUnknown, "", "", "foobar")
		c.Assert(ipAddr.Life(), gc.Equals, state.Alive)

		// verify the address was stored in the state
		ipAddr, err = s.State.IPAddress(test)
		c.Assert(err, jc.ErrorIsNil)
		s.assertAddress(c, ipAddr, addr, state.AddressStateUnknown, "", "", "foobar")
	}
}

func (s *IPAddressSuite) TestAddIPAddressInvalid(c *gc.C) {
	addr := network.Address{Value: "foo"}
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.Satisfies, errors.IsNotValid)
	c.Assert(err, gc.ErrorMatches, `cannot add IP address "foo": address not valid`)
}

func (s *IPAddressSuite) TestAddIPAddressAlreadyExists(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	_, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	_, err = s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.Satisfies, errors.IsAlreadyExists)
	c.Assert(err, gc.ErrorMatches,
		`cannot add IP address "public:0.1.2.3": address already exists`,
	)
}

func (s *IPAddressSuite) TestIPAddressNotFound(c *gc.C) {
	_, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
	c.Assert(err, gc.ErrorMatches, `IP address "0.1.2.3" not found`)
}

func (s *IPAddressSuite) TestEnsureDeadRemove(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	// Should not be able to remove an Alive IP address.
	c.Assert(ipAddr.Life(), gc.Equals, state.Alive)
	err = ipAddr.Remove()
	msg := fmt.Sprintf("cannot remove IP address %q: IP address is not dead", ipAddr.String())
	c.Assert(err, gc.ErrorMatches, msg)

	err = ipAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	// EnsureDead twice should not be an error
	err = ipAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	// Remove twice is also fine.
	err = ipAddr.Remove()
	c.Assert(err, jc.ErrorIsNil)

	_, err = s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.Satisfies, errors.IsNotFound)
	c.Assert(err, gc.ErrorMatches, `IP address "0.1.2.3" not found`)
}

func (s *IPAddressSuite) TestSetStateDead(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	copyIPAddr, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	err = copyIPAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	err = ipAddr.SetState(state.AddressStateAllocated)
	msg := fmt.Sprintf(`cannot set IP address %q to state "allocated": address is dead`, ipAddr.String())
	c.Assert(err, gc.ErrorMatches, msg)
}

func (s *IPAddressSuite) TestAllocateToDead(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	copyIPAddr, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	err = copyIPAddr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	msg := fmt.Sprintf(`cannot allocate IP address %q to machine "foobar", interface "wibble": address is dead`, ipAddr.String())
	err = ipAddr.AllocateTo("foobar", "wibble")
	c.Assert(err, gc.ErrorMatches, msg)
}

func (s *IPAddressSuite) TestAddressStateString(c *gc.C) {
	for i, test := range []struct {
		ipState state.AddressState
		expect  string
	}{{
		state.AddressStateUnknown,
		"<unknown>",
	}, {
		state.AddressStateAllocated,
		"allocated",
	}, {
		state.AddressStateUnavailable,
		"unavailable",
	}} {
		c.Logf("test %d: %q -> %q", i, test.ipState, test.expect)
		c.Check(test.ipState.String(), gc.Equals, test.expect)
	}
}

func (s *IPAddressSuite) TestSetState(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)

	for i, test := range []struct {
		initial, changeTo state.AddressState
		err               string
	}{{
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateUnknown,
	}, {
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateAllocated,
	}, {
		initial:  state.AddressStateUnknown,
		changeTo: state.AddressStateUnavailable,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateAllocated,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateUnavailable,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateUnknown,
		err: `cannot set IP address "public:0.1.2.3" to state "<unknown>": ` +
			`transition from "allocated" not valid`,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateUnknown,
		err: `cannot set IP address "public:0.1.2.3" to state "<unknown>": ` +
			`transition from "unavailable" not valid`,
	}, {
		initial:  state.AddressStateAllocated,
		changeTo: state.AddressStateUnavailable,
		err: `cannot set IP address "public:0.1.2.3" to state "unavailable": ` +
			`transition from "allocated" not valid`,
	}, {
		initial:  state.AddressStateUnavailable,
		changeTo: state.AddressStateAllocated,
		err: `cannot set IP address "public:0.1.2.3" to state "allocated": ` +
			`transition from "unavailable" not valid`,
	}} {
		c.Logf("test %d: %q -> %q ok:%v", i, test.initial, test.changeTo, test.err == "")
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Check(err, jc.ErrorIsNil)

		// Initially, all addresses have AddressStateUnknown.
		c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)

		if test.initial != state.AddressStateUnknown {
			err = ipAddr.SetState(test.initial)
			c.Check(err, jc.ErrorIsNil)
		}
		err = ipAddr.SetState(test.changeTo)
		if test.err != "" {
			c.Check(err, gc.ErrorMatches, test.err)
			c.Check(err, jc.Satisfies, errors.IsNotValid)
			c.Check(ipAddr.EnsureDead(), jc.ErrorIsNil)
			c.Check(ipAddr.Remove(), jc.ErrorIsNil)
			continue
		}
		c.Check(err, jc.ErrorIsNil)
		c.Check(ipAddr.State(), gc.Equals, test.changeTo)
		c.Check(ipAddr.EnsureDead(), jc.ErrorIsNil)
		c.Check(ipAddr.Remove(), jc.ErrorIsNil)
	}
}

func (s *IPAddressSuite) TestAllocateTo(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateUnknown)
	c.Assert(ipAddr.MachineId(), gc.Equals, "")
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "")

	err = ipAddr.AllocateTo("wibble", "wobble")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.State(), gc.Equals, state.AddressStateAllocated)
	c.Assert(ipAddr.MachineId(), gc.Equals, "wibble")
	c.Assert(ipAddr.InterfaceId(), gc.Equals, "wobble")

	freshCopy, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(freshCopy.State(), gc.Equals, state.AddressStateAllocated)
	c.Assert(freshCopy.MachineId(), gc.Equals, "wibble")
	c.Assert(freshCopy.InterfaceId(), gc.Equals, "wobble")

	// allocating twice should fail.
	err = ipAddr.AllocateTo("m", "i")
	c.Assert(err, gc.ErrorMatches,
		`cannot allocate IP address "public:0.1.2.3" to machine "m", interface "i": `+
			`already allocated or unavailable`,
	)
}

func (s *IPAddressSuite) TestAddress(c *gc.C) {
	addr := network.NewAddress("0.1.2.3", network.ScopePublic)
	ipAddr, err := s.State.AddIPAddress(addr, "foobar")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddr.Address(), jc.DeepEquals, addr)

}

func (s *IPAddressSuite) TestAllocatedIPAddresses(c *gc.C) {
	addresses := [][]string{
		{"0.1.2.3", "wibble"},
		{"0.1.2.4", "wibble"},
		{"0.1.2.5", "wobble"},
	}
	for _, details := range addresses {
		addr := network.NewAddress(details[0], network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		err = ipAddr.AllocateTo(details[1], "wobble")
		c.Assert(err, jc.ErrorIsNil)
	}
	result, err := s.State.AllocatedIPAddresses("wibble")
	c.Assert(err, jc.ErrorIsNil)
	addr1, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	addr2, err := s.State.IPAddress("0.1.2.4")
	c.Assert(err, jc.ErrorIsNil)
	expected := []*state.IPAddress{addr1, addr2}
	c.Assert(result, jc.SameContents, expected)

}

func (s *IPAddressSuite) TestDeadIPAddresses(c *gc.C) {
	addresses := [][]string{
		{"0.1.2.3", "wibble"},
		{"0.1.2.4", "wibble"},
		{"0.1.2.5", "wobble"},
		{"0.1.2.6", "wobble"},
	}
	for i, details := range addresses {
		addr := network.NewAddress(details[0], network.ScopePublic)
		ipAddr, err := s.State.AddIPAddress(addr, "foobar")
		c.Assert(err, jc.ErrorIsNil)
		err = ipAddr.AllocateTo(details[1], "wobble")
		c.Assert(err, jc.ErrorIsNil)
		if i%2 == 0 {
			err := ipAddr.EnsureDead()
			c.Assert(err, jc.ErrorIsNil)
		} else {
			c.Assert(ipAddr.Life(), gc.Equals, state.Alive)
		}
	}

	ipAddresses, err := s.State.DeadIPAddresses()
	c.Assert(err, jc.ErrorIsNil)
	addr1, err := s.State.IPAddress("0.1.2.3")
	c.Assert(err, jc.ErrorIsNil)
	addr3, err := s.State.IPAddress("0.1.2.5")
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(ipAddresses, jc.SameContents, []*state.IPAddress{addr1, addr3})
}

func (s *IPAddressSuite) TestRefresh(c *gc.C) {
	rawAddr := network.NewAddress("0.1.2.3", network.ScopeUnknown)
	addr, err := s.State.AddIPAddress(rawAddr, "foobar")
	c.Assert(err, jc.ErrorIsNil)

	addrCopy, err := s.State.IPAddress(rawAddr.Value)
	c.Assert(err, jc.ErrorIsNil)

	err = addr.EnsureDead()
	c.Assert(err, jc.ErrorIsNil)

	c.Assert(addrCopy.Life(), gc.Equals, state.Alive)
	err = addrCopy.Refresh()
	c.Assert(err, jc.ErrorIsNil)
	c.Assert(addrCopy.Life(), gc.Equals, state.Dead)
}
<|endoftext|>


<commit_before>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 builder

import (
	"fmt"

	"github.com/spf13/cobra"
	"mynewt.apache.org/newt/newt/cli"
	"mynewt.apache.org/newt/newt/pkg"
	"mynewt.apache.org/newt/newt/project"
	"mynewt.apache.org/newt/newt/target"
	"mynewt.apache.org/newt/util"
)

const TARGET_TEST_NAME = "targets/unittest"

func buildRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	targetName := args[0]
	t := target.GetTargets()[targetName]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid target name"))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}
	err = b.Build()
	if err != nil {
		fmt.Println(err)
	}
}

func cleanRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	targetName := args[0]
	t := target.GetTargets()[targetName]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid target name"))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}
	err = b.Clean()
	if err != nil {
		fmt.Println(err)
	}
}

func testRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	t := target.GetTargets()[TARGET_TEST_NAME]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("can't find unit test target: "+
			TARGET_TEST_NAME))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}

	pkgString := args[0]
	dep, err := pkg.NewDependency(nil, pkgString)
	if err != nil {
		cli.NewtUsage(cmd, err)
	}
	if dep == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid package name"))
	}
	pack := project.GetProject().ResolveDependency(dep)
	if pack == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid package name"))
	}

	err = b.Test(pack.(*pkg.LocalPackage))
	if err != nil {
		fmt.Println(err)
	}
}

func AddCommands(cmd *cobra.Command) {
	buildHelpText := ""
	buildHelpEx := ""
	buildCmd := &cobra.Command{
		Use:     "build",
		Short:   "Commands for building targets",
		Long:    buildHelpText,
		Example: buildHelpEx,
		Run:     buildRunCmd,
	}

	cmd.AddCommand(buildCmd)

	cleanHelpText := ""
	cleanHelpEx := ""
	cleanCmd := &cobra.Command{
		Use:     "clean",
		Short:   "Commands for cleaning targets",
		Long:    cleanHelpText,
		Example: cleanHelpEx,
		Run:     cleanRunCmd,
	}

	cmd.AddCommand(cleanCmd)

	testHelpText := ""
	testHelpEx := ""
	testCmd := &cobra.Command{
		Use:     "test",
		Short:   "Executes unit tests for one or more packages",
		Long:    testHelpText,
		Example: testHelpEx,
		Run:     testRunCmd,
	}

	cmd.AddCommand(testCmd)
}
<commit_msg>Allow multiple packages for test command.<commit_after>/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 builder

import (
	"fmt"

	"github.com/spf13/cobra"
	"mynewt.apache.org/newt/newt/cli"
	"mynewt.apache.org/newt/newt/pkg"
	"mynewt.apache.org/newt/newt/project"
	"mynewt.apache.org/newt/newt/target"
	"mynewt.apache.org/newt/util"
)

const TARGET_TEST_NAME = "targets/unittest"

func buildRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	targetName := args[0]
	t := target.GetTargets()[targetName]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid target name"))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}
	err = b.Build()
	if err != nil {
		fmt.Println(err)
	}
}

func cleanRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	targetName := args[0]
	t := target.GetTargets()[targetName]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("invalid target name"))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}
	err = b.Clean()
	if err != nil {
		fmt.Println(err)
	}
}

func testRunCmd(cmd *cobra.Command, args []string) {
	if len(args) < 1 {
		cli.NewtUsage(cmd, nil)
	}

	t := target.GetTargets()[TARGET_TEST_NAME]
	if t == nil {
		cli.NewtUsage(cmd, util.NewNewtError("can't find unit test target: "+
			TARGET_TEST_NAME))
	}

	b, err := NewBuilder(t)
	if err != nil {
		fmt.Println(err)
	}

	// Verify and resolve each specified package.
	packs := []*pkg.LocalPackage{}
	for _, pkgName := range args {
		dep, err := pkg.NewDependency(nil, pkgName)
		if err != nil {
			cli.NewtUsage(cmd, err)
		}
		if dep == nil {
			cli.NewtUsage(cmd, util.NewNewtError("invalid package name: "+
				pkgName))
		}
		pack := project.GetProject().ResolveDependency(dep)
		if pack == nil {
			cli.NewtUsage(cmd, util.NewNewtError("unknown package: "+pkgName))
		}

		packs = append(packs, pack.(*pkg.LocalPackage))
	}

	// Test each package.
	for _, pack := range packs {
		err = b.Test(pack)
		if err != nil {
			fmt.Println(err)
		}
	}
}

func AddCommands(cmd *cobra.Command) {
	buildHelpText := ""
	buildHelpEx := ""
	buildCmd := &cobra.Command{
		Use:     "build",
		Short:   "Commands for building targets",
		Long:    buildHelpText,
		Example: buildHelpEx,
		Run:     buildRunCmd,
	}

	cmd.AddCommand(buildCmd)

	cleanHelpText := ""
	cleanHelpEx := ""
	cleanCmd := &cobra.Command{
		Use:     "clean",
		Short:   "Commands for cleaning targets",
		Long:    cleanHelpText,
		Example: cleanHelpEx,
		Run:     cleanRunCmd,
	}

	cmd.AddCommand(cleanCmd)

	testHelpText := ""
	testHelpEx := ""
	testCmd := &cobra.Command{
		Use:     "test",
		Short:   "Executes unit tests for one or more packages",
		Long:    testHelpText,
		Example: testHelpEx,
		Run:     testRunCmd,
	}

	cmd.AddCommand(testCmd)
}
<|endoftext|>


<commit_before>package ma

import (
	"yap/alg/graph"
	"yap/nlp/format/lex"
	. "yap/nlp/types"
	"yap/util"

	"fmt"
	"log"
)

const ESTIMATED_MORPHS_PER_TOKEN = 5

type BGULex struct {
	MaxPrefixLen int
	Prefixes     map[string][]BasicMorphemes

	Lex map[string][]BasicMorphemes

	Files []string
	Stats *AnalyzeStats
}

var _ MorphologicalAnalyzer = &BGULex{}

func (l *BGULex) loadTokens(file, format string) {
	tokens, err := lex.ReadFile(file, format)
	if err != nil {
		panic(fmt.Sprintf("Failed to load %v: %v", file, err))
	}
	var m map[string][]BasicMorphemes
	if format == "prefix" {
		l.Prefixes = make(map[string][]BasicMorphemes, len(tokens))
		m = l.Prefixes
	} else if format == "lexicon" {
		l.Lex = make(map[string][]BasicMorphemes, len(tokens))
		m = l.Lex
	}
	log.Println("Found", len(tokens), "tokens in lexicon file:", file)
	for _, token := range tokens {
		if cur, exists := m[token.Token]; exists {
			m[token.Token] = append(cur, token.Morphemes...)
		} else {
			m[token.Token] = token.Morphemes
		}
	}
}

func (l *BGULex) LoadPrefixes(file string) {
	l.loadTokens(file, "prefix")
	l.MaxPrefixLen = 0
	for _, morphs := range l.Prefixes {
		if l.MaxPrefixLen < len(morphs) {
			l.MaxPrefixLen = len(morphs)
		}
	}
	log.Println("Loaded", len(l.Prefixes), "prefixes from lexicon")
}

func (l *BGULex) LoadLex(file string) {
	l.loadTokens(file, "lexicon")
	log.Println("Loaded", len(l.Lex), "tokens from lexicon")
}

func (l *BGULex) OOVAnalysis(input string) []BasicMorphemes {
	return []BasicMorphemes{BasicMorphemes([]*Morpheme{
		&Morpheme{
			BasicDirectedEdge: graph.BasicDirectedEdge{0, 0, 1},
			Form:              input,
			Lemma:             "",
			CPOS:              "NNP",
			POS:               "NNP",
			FeatureStr:        "",
		},
	})}
}
func (l *BGULex) AnalyzeToken(input string, startingNode, numToken int) (*Lattice, interface{}) {
	// log.Println("Analyzing token", numToken, "starting at", startingNode)
	lat := &Lattice{
		Token:     Token(input),
		Morphemes: make(Morphemes, 0, ESTIMATED_MORPHS_PER_TOKEN),
		Next:      make(map[int][]int, ESTIMATED_MORPHS_PER_TOKEN),
		BottomId:  startingNode,
		TopId:     startingNode,
	}
	lat.Next[0] = make([]int, 0, 1)
	var (
		prefixLat, hostLat                  []BasicMorphemes
		prefixExists, hostExists, anyExists bool
	)
	hostLat, hostExists = l.Lex[input]
	if hostExists {
		// log.Println("\tPrefix 0")
		lat.AddAnalysis(nil, hostLat, numToken)
		anyExists = true
	}
	for i := 1; i < util.Min(l.MaxPrefixLen, len(input)); i++ {
		prefixLat, prefixExists = l.Prefixes[input[:i]]
		if prefixExists {
			hostLat, hostExists = l.Lex[input[i:]]
			if hostExists {
				// log.Println("\tPrefix", i)
				lat.AddAnalysis(prefixLat, hostLat, numToken)
				anyExists = true
			}
		}
	}
	if !anyExists {
		log.Println("Token is OOV:", input)
		if l.Stats != nil {
			l.Stats.OOVTokens++
			l.Stats.AddOOVToken(input)
		}
		hostLat = l.OOVAnalysis(input)
		lat.AddAnalysis(nil, hostLat, numToken)
	}
	lat.Optimize()
	return lat, nil
}

func (l *BGULex) Analyze(input []string) (LatticeSentence, interface{}) {
	retval := make(LatticeSentence, len(input))
	var (
		lat     *Lattice
		curNode int
	)
	for i, token := range input {
		if l.Stats != nil {
			l.Stats.TotalTokens++
			l.Stats.AddToken(token)
		}
		lat, _ = l.AnalyzeToken(token, curNode, i)
		curNode = lat.Top()
		// log.Println("New top is", curNode)
		retval[i] = *lat
	}
	return retval, nil
}
<commit_msg>Refactor and fix prefix bug<commit_after>package ma

import (
	"yap/alg/graph"
	"yap/nlp/format/lex"
	. "yap/nlp/types"
	"yap/util"

	"fmt"
	"log"
)

const ESTIMATED_MORPHS_PER_TOKEN = 5

type BGULex struct {
	MaxPrefixLen int
	Prefixes     map[string][]BasicMorphemes

	Lex map[string][]BasicMorphemes

	Files []string
	Stats *AnalyzeStats
}

var _ MorphologicalAnalyzer = &BGULex{}

func (l *BGULex) loadTokens(file, format string) {
	tokens, err := lex.ReadFile(file, format)
	if err != nil {
		panic(fmt.Sprintf("Failed to load %v: %v", file, err))
	}
	var m map[string][]BasicMorphemes
	if format == "prefix" {
		l.Prefixes = make(map[string][]BasicMorphemes, len(tokens))
		m = l.Prefixes
	} else if format == "lexicon" {
		l.Lex = make(map[string][]BasicMorphemes, len(tokens))
		m = l.Lex
	}
	log.Println("Found", len(tokens), "tokens in lexicon file:", file)
	for _, token := range tokens {
		if cur, exists := m[token.Token]; exists {
			m[token.Token] = append(cur, token.Morphemes...)
		} else {
			m[token.Token] = token.Morphemes
		}
	}
}

func (l *BGULex) LoadPrefixes(file string) {
	l.loadTokens(file, "prefix")
	l.MaxPrefixLen = 0
	for _, morphs := range l.Prefixes {
		if l.MaxPrefixLen < len(morphs) {
			l.MaxPrefixLen = len(morphs)
		}
	}
	log.Println("Loaded", len(l.Prefixes), "prefixes from lexicon")
}

func (l *BGULex) LoadLex(file string) {
	l.loadTokens(file, "lexicon")
	log.Println("Loaded", len(l.Lex), "tokens from lexicon")
}

func (l *BGULex) OOVAnalysis(input string) []BasicMorphemes {
	return []BasicMorphemes{BasicMorphemes([]*Morpheme{
		&Morpheme{
			BasicDirectedEdge: graph.BasicDirectedEdge{0, 0, 1},
			Form:              input,
			Lemma:             "",
			CPOS:              "NNP",
			POS:               "NNP",
			FeatureStr:        "",
		},
	})}
}

func (l *BGULex) analyzeTokenForLen(lat *Lattice, input string, startingNode, numToken, prefixLen int) bool {
	var found bool
	if len(input) < prefixLen*2 {
		return found
	}
	prefixLat, prefixExists := l.Prefixes[input[0:prefixLen*2]]
	log.Println("\tPrefix", input[0:prefixLen*2], prefixExists)
	if prefixExists {
		hostLat, hostExists := l.Lex[input[2*prefixLen:]]
		log.Println("\t\tHost", input[2*prefixLen:], hostExists)
		if hostExists {
			log.Println("\t\tAdding", prefixLat, hostLat)
			lat.AddAnalysis(prefixLat, hostLat, numToken)
			found = true
		}
	}
	return found
}
func (l *BGULex) AnalyzeToken(input string, startingNode, numToken int) (*Lattice, interface{}) {
	log.Println("Analyzing token", numToken, "starting at", startingNode)
	lat := &Lattice{
		Token:     Token(input),
		Morphemes: make(Morphemes, 0, ESTIMATED_MORPHS_PER_TOKEN),
		Next:      make(map[int][]int, ESTIMATED_MORPHS_PER_TOKEN),
		BottomId:  startingNode,
		TopId:     startingNode,
	}
	lat.Next[0] = make([]int, 0, 1)
	var (
		hostLat               []BasicMorphemes
		hostExists, anyExists bool
	)
	hostLat, hostExists = l.Lex[input]
	if hostExists {
		log.Println("\tPrefix 0")
		lat.AddAnalysis(nil, hostLat, numToken)
		anyExists = true
	}
	for i := 1; i < util.Min(l.MaxPrefixLen, len(input)); i++ {
		log.Println("\ti is", i)
		found := l.analyzeTokenForLen(lat, input, startingNode, numToken, i)
		anyExists = anyExists || found
	}
	if !anyExists {
		// log.Println("Token is OOV:", input)
		if l.Stats != nil {
			l.Stats.OOVTokens++
			l.Stats.AddOOVToken(input)
		}
		hostLat = l.OOVAnalysis(input)
		lat.AddAnalysis(nil, hostLat, numToken)
	}
	lat.Optimize()
	return lat, nil
}

func (l *BGULex) Analyze(input []string) (LatticeSentence, interface{}) {
	retval := make(LatticeSentence, len(input))
	var (
		lat     *Lattice
		curNode int
	)
	for i, token := range input {
		if l.Stats != nil {
			l.Stats.TotalTokens++
			l.Stats.AddToken(token)
		}
		lat, _ = l.AnalyzeToken(token, curNode, i)
		curNode = lat.Top()
		// log.Println("New top is", curNode)
		retval[i] = *lat
	}
	return retval, nil
}
<|endoftext|>