<commit_before>package sarama

import (
	"fmt"
)

// TestReporter has methods matching go's testing.T to avoid importing
// `testing` in the main part of the library.
type TestReporter interface {
	Error(...interface{})
	Errorf(string, ...interface{})
	Fatal(...interface{})
	Fatalf(string, ...interface{})
}

// MockResponse is a response builder interface it defines one method that
// allows generating a response based on a request body. MockResponses are used
// to program behavior of MockBroker in tests.
type MockResponse interface {
	For(reqBody versionedDecoder) (res encoder)
}

// MockWrapper is a mock response builder that returns a particular concrete
// response regardless of the actual request passed to the `For` method.
type MockWrapper struct {
	res encoder
}

func (mw *MockWrapper) For(reqBody versionedDecoder) (res encoder) {
	return mw.res
}

func NewMockWrapper(res encoder) *MockWrapper {
	return &MockWrapper{res: res}
}

// MockSequence is a mock response builder that is created from a sequence of
// concrete responses. Every time when a `MockBroker` calls its `For` method
// the next response from the sequence is returned. When the end of the
// sequence is reached the last element from the sequence is returned.
type MockSequence struct {
	responses []MockResponse
}

func NewMockSequence(responses ...interface{}) *MockSequence {
	ms := &MockSequence{}
	ms.responses = make([]MockResponse, len(responses))
	for i, res := range responses {
		switch res := res.(type) {
		case MockResponse:
			ms.responses[i] = res
		case encoder:
			ms.responses[i] = NewMockWrapper(res)
		default:
			panic(fmt.Sprintf("Unexpected response type: %T", res))
		}
	}
	return ms
}

func (mc *MockSequence) For(reqBody versionedDecoder) (res encoder) {
	res = mc.responses[0].For(reqBody)
	if len(mc.responses) > 1 {
		mc.responses = mc.responses[1:]
	}
	return res
}

// MockMetadataResponse is a `MetadataResponse` builder.
type MockMetadataResponse struct {
	leaders map[string]map[int32]int32
	brokers map[string]int32
	t       TestReporter
}

func NewMockMetadataResponse(t TestReporter) *MockMetadataResponse {
	return &MockMetadataResponse{
		leaders: make(map[string]map[int32]int32),
		brokers: make(map[string]int32),
		t:       t,
	}
}

func (mmr *MockMetadataResponse) SetLeader(topic string, partition, brokerID int32) *MockMetadataResponse {
	partitions := mmr.leaders[topic]
	if partitions == nil {
		partitions = make(map[int32]int32)
		mmr.leaders[topic] = partitions
	}
	partitions[partition] = brokerID
	return mmr
}

func (mmr *MockMetadataResponse) SetBroker(addr string, brokerID int32) *MockMetadataResponse {
	mmr.brokers[addr] = brokerID
	return mmr
}

func (mmr *MockMetadataResponse) For(reqBody versionedDecoder) encoder {
	metadataRequest := reqBody.(*MetadataRequest)
	metadataResponse := &MetadataResponse{}
	for addr, brokerID := range mmr.brokers {
		metadataResponse.AddBroker(addr, brokerID)
	}
	if len(metadataRequest.Topics) == 0 {
		for topic, partitions := range mmr.leaders {
			for partition, brokerID := range partitions {
				metadataResponse.AddTopicPartition(topic, partition, brokerID, nil, nil, ErrNoError)
			}
		}
		return metadataResponse
	}
	for _, topic := range metadataRequest.Topics {
		for partition, brokerID := range mmr.leaders[topic] {
			metadataResponse.AddTopicPartition(topic, partition, brokerID, nil, nil, ErrNoError)
		}
	}
	return metadataResponse
}

// MockOffsetResponse is an `OffsetResponse` builder.
type MockOffsetResponse struct {
	offsets map[string]map[int32]map[int64]int64
	t       TestReporter
}

func NewMockOffsetResponse(t TestReporter) *MockOffsetResponse {
	return &MockOffsetResponse{
		offsets: make(map[string]map[int32]map[int64]int64),
		t:       t,
	}
}

func (mor *MockOffsetResponse) SetOffset(topic string, partition int32, time, offset int64) *MockOffsetResponse {
	partitions := mor.offsets[topic]
	if partitions == nil {
		partitions = make(map[int32]map[int64]int64)
		mor.offsets[topic] = partitions
	}
	times := partitions[partition]
	if times == nil {
		times = make(map[int64]int64)
		partitions[partition] = times
	}
	times[time] = offset
	return mor
}

func (mor *MockOffsetResponse) For(reqBody versionedDecoder) encoder {
	offsetRequest := reqBody.(*OffsetRequest)
	offsetResponse := &OffsetResponse{}
	for topic, partitions := range offsetRequest.blocks {
		for partition, block := range partitions {
			offset := mor.getOffset(topic, partition, block.time)
			offsetResponse.AddTopicPartition(topic, partition, offset)
		}
	}
	return offsetResponse
}

func (mor *MockOffsetResponse) getOffset(topic string, partition int32, time int64) int64 {
	partitions := mor.offsets[topic]
	if partitions == nil {
		mor.t.Errorf("missing topic: %s", topic)
	}
	times := partitions[partition]
	if times == nil {
		mor.t.Errorf("missing partition: %d", partition)
	}
	offset, ok := times[time]
	if !ok {
		mor.t.Errorf("missing time: %d", time)
	}
	return offset
}

// MockFetchResponse is a `FetchResponse` builder.
type MockFetchResponse struct {
	messages       map[string]map[int32]map[int64]Encoder
	highWaterMarks map[string]map[int32]int64
	t              TestReporter
	batchSize      int
}

func NewMockFetchResponse(t TestReporter, batchSize int) *MockFetchResponse {
	return &MockFetchResponse{
		messages:       make(map[string]map[int32]map[int64]Encoder),
		highWaterMarks: make(map[string]map[int32]int64),
		t:              t,
		batchSize:      batchSize,
	}
}

func (mfr *MockFetchResponse) SetMessage(topic string, partition int32, offset int64, msg Encoder) *MockFetchResponse {
	partitions := mfr.messages[topic]
	if partitions == nil {
		partitions = make(map[int32]map[int64]Encoder)
		mfr.messages[topic] = partitions
	}
	messages := partitions[partition]
	if messages == nil {
		messages = make(map[int64]Encoder)
		partitions[partition] = messages
	}
	messages[offset] = msg
	return mfr
}

func (mfr *MockFetchResponse) SetHighWaterMark(topic string, partition int32, offset int64) *MockFetchResponse {
	partitions := mfr.highWaterMarks[topic]
	if partitions == nil {
		partitions = make(map[int32]int64)
		mfr.highWaterMarks[topic] = partitions
	}
	partitions[partition] = offset
	return mfr
}

func (mfr *MockFetchResponse) For(reqBody versionedDecoder) encoder {
	fetchRequest := reqBody.(*FetchRequest)
	res := &FetchResponse{}
	for topic, partitions := range fetchRequest.blocks {
		for partition, block := range partitions {
			initialOffset := block.fetchOffset
			offset := initialOffset
			maxOffset := initialOffset + int64(mfr.getMessageCount(topic, partition))
			for i := 0; i < mfr.batchSize && offset < maxOffset; {
				msg := mfr.getMessage(topic, partition, offset)
				if msg != nil {
					res.AddMessage(topic, partition, nil, msg, offset)
					i++
				}
				offset++
			}
			fb := res.GetBlock(topic, partition)
			if fb == nil {
				res.AddError(topic, partition, ErrNoError)
				fb = res.GetBlock(topic, partition)
			}
			fb.HighWaterMarkOffset = mfr.getHighWaterMark(topic, partition)
		}
	}
	return res
}

func (mfr *MockFetchResponse) getMessage(topic string, partition int32, offset int64) Encoder {
	partitions := mfr.messages[topic]
	if partitions == nil {
		return nil
	}
	messages := partitions[partition]
	if messages == nil {
		return nil
	}
	return messages[offset]
}

func (mfr *MockFetchResponse) getMessageCount(topic string, partition int32) int {
	partitions := mfr.messages[topic]
	if partitions == nil {
		return 0
	}
	messages := partitions[partition]
	if messages == nil {
		return 0
	}
	return len(messages)
}

func (mfr *MockFetchResponse) getHighWaterMark(topic string, partition int32) int64 {
	partitions := mfr.highWaterMarks[topic]
	if partitions == nil {
		return 0
	}
	return partitions[partition]
}

// MockConsumerMetadataResponse is a `ConsumerMetadataResponse` builder.
type MockConsumerMetadataResponse struct {
	coordinators map[string]interface{}
	t            TestReporter
}

func NewMockConsumerMetadataResponse(t TestReporter) *MockConsumerMetadataResponse {
	return &MockConsumerMetadataResponse{
		coordinators: make(map[string]interface{}),
		t:            t,
	}
}

func (mr *MockConsumerMetadataResponse) SetCoordinator(group string, broker *MockBroker) *MockConsumerMetadataResponse {
	mr.coordinators[group] = broker
	return mr
}

func (mr *MockConsumerMetadataResponse) SetError(group string, kerror KError) *MockConsumerMetadataResponse {
	mr.coordinators[group] = kerror
	return mr
}

func (mr *MockConsumerMetadataResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*ConsumerMetadataRequest)
	group := req.ConsumerGroup
	res := &ConsumerMetadataResponse{}
	v := mr.coordinators[group]
	switch v := v.(type) {
	case *MockBroker:
		res.Coordinator = &Broker{id: v.BrokerID(), addr: v.Addr()}
	case KError:
		res.Err = v
	}
	return res
}

// MockOffsetCommitResponse is a `OffsetCommitResponse` builder.
type MockOffsetCommitResponse struct {
	errors map[string]map[string]map[int32]KError
	t      TestReporter
}

func NewMockOffsetCommitResponse(t TestReporter) *MockOffsetCommitResponse {
	return &MockOffsetCommitResponse{t: t}
}

func (mr *MockOffsetCommitResponse) SetError(group, topic string, partition int32, kerror KError) *MockOffsetCommitResponse {
	if mr.errors == nil {
		mr.errors = make(map[string]map[string]map[int32]KError)
	}
	topics := mr.errors[group]
	if topics == nil {
		topics = make(map[string]map[int32]KError)
		mr.errors[group] = topics
	}
	partitions := topics[topic]
	if partitions == nil {
		partitions = make(map[int32]KError)
		topics[topic] = partitions
	}
	partitions[partition] = kerror
	return mr
}

func (mr *MockOffsetCommitResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*OffsetCommitRequest)
	group := req.ConsumerGroup
	res := &OffsetCommitResponse{}
	for topic, partitions := range req.blocks {
		for partition := range partitions {
			res.AddError(topic, partition, mr.getError(group, topic, partition))
		}
	}
	return res
}

func (mr *MockOffsetCommitResponse) getError(group, topic string, partition int32) KError {
	topics := mr.errors[group]
	if topics == nil {
		return ErrNoError
	}
	partitions := topics[topic]
	if partitions == nil {
		return ErrNoError
	}
	kerror, ok := partitions[partition]
	if !ok {
		return ErrNoError
	}
	return kerror
}

// MockProduceResponse is a `ProduceResponse` builder.
type MockProduceResponse struct {
	errors map[string]map[int32]KError
	t      TestReporter
}

func NewMockProduceResponse(t TestReporter) *MockProduceResponse {
	return &MockProduceResponse{t: t}
}

func (mr *MockProduceResponse) SetError(topic string, partition int32, kerror KError) *MockProduceResponse {
	if mr.errors == nil {
		mr.errors = make(map[string]map[int32]KError)
	}
	partitions := mr.errors[topic]
	if partitions == nil {
		partitions = make(map[int32]KError)
		mr.errors[topic] = partitions
	}
	partitions[partition] = kerror
	return mr
}

func (mr *MockProduceResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*ProduceRequest)
	res := &ProduceResponse{}
	for topic, partitions := range req.msgSets {
		for partition := range partitions {
			res.AddTopicPartition(topic, partition, mr.getError(topic, partition))
		}
	}
	return res
}

func (mr *MockProduceResponse) getError(topic string, partition int32) KError {
	partitions := mr.errors[topic]
	if partitions == nil {
		return ErrNoError
	}
	kerror, ok := partitions[partition]
	if !ok {
		return ErrNoError
	}
	return kerror
}

// MockOffsetFetchResponse is a `OffsetFetchResponse` builder.
type MockOffsetFetchResponse struct {
	offsets map[string]map[string]map[int32]*OffsetFetchResponseBlock
	t       TestReporter
}

func NewMockOffsetFetchResponse(t TestReporter) *MockOffsetFetchResponse {
	return &MockOffsetFetchResponse{t: t}
}

func (mr *MockOffsetFetchResponse) SetOffset(group, topic string, partition int32, offset int64, metadata string, kerror KError) *MockOffsetFetchResponse {
	if mr.offsets == nil {
		mr.offsets = make(map[string]map[string]map[int32]*OffsetFetchResponseBlock)
	}
	topics := mr.offsets[group]
	if topics == nil {
		topics = make(map[string]map[int32]*OffsetFetchResponseBlock)
		mr.offsets[group] = topics
	}
	partitions := topics[topic]
	if partitions == nil {
		partitions = make(map[int32]*OffsetFetchResponseBlock)
		topics[topic] = partitions
	}
	partitions[partition] = &OffsetFetchResponseBlock{offset, metadata, kerror}
	return mr
}

func (mr *MockOffsetFetchResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*OffsetFetchRequest)
	group := req.ConsumerGroup
	res := &OffsetFetchResponse{}
	for topic, partitions := range mr.offsets[group] {
		for partition, block := range partitions {
			res.AddBlock(topic, partition, block)
		}
	}
	return res
}
<commit_msg>Permit setting version on mock fetch response<commit_after>package sarama

import (
	"fmt"
)

// TestReporter has methods matching go's testing.T to avoid importing
// `testing` in the main part of the library.
type TestReporter interface {
	Error(...interface{})
	Errorf(string, ...interface{})
	Fatal(...interface{})
	Fatalf(string, ...interface{})
}

// MockResponse is a response builder interface it defines one method that
// allows generating a response based on a request body. MockResponses are used
// to program behavior of MockBroker in tests.
type MockResponse interface {
	For(reqBody versionedDecoder) (res encoder)
}

// MockWrapper is a mock response builder that returns a particular concrete
// response regardless of the actual request passed to the `For` method.
type MockWrapper struct {
	res encoder
}

func (mw *MockWrapper) For(reqBody versionedDecoder) (res encoder) {
	return mw.res
}

func NewMockWrapper(res encoder) *MockWrapper {
	return &MockWrapper{res: res}
}

// MockSequence is a mock response builder that is created from a sequence of
// concrete responses. Every time when a `MockBroker` calls its `For` method
// the next response from the sequence is returned. When the end of the
// sequence is reached the last element from the sequence is returned.
type MockSequence struct {
	responses []MockResponse
}

func NewMockSequence(responses ...interface{}) *MockSequence {
	ms := &MockSequence{}
	ms.responses = make([]MockResponse, len(responses))
	for i, res := range responses {
		switch res := res.(type) {
		case MockResponse:
			ms.responses[i] = res
		case encoder:
			ms.responses[i] = NewMockWrapper(res)
		default:
			panic(fmt.Sprintf("Unexpected response type: %T", res))
		}
	}
	return ms
}

func (mc *MockSequence) For(reqBody versionedDecoder) (res encoder) {
	res = mc.responses[0].For(reqBody)
	if len(mc.responses) > 1 {
		mc.responses = mc.responses[1:]
	}
	return res
}

// MockMetadataResponse is a `MetadataResponse` builder.
type MockMetadataResponse struct {
	leaders map[string]map[int32]int32
	brokers map[string]int32
	t       TestReporter
}

func NewMockMetadataResponse(t TestReporter) *MockMetadataResponse {
	return &MockMetadataResponse{
		leaders: make(map[string]map[int32]int32),
		brokers: make(map[string]int32),
		t:       t,
	}
}

func (mmr *MockMetadataResponse) SetLeader(topic string, partition, brokerID int32) *MockMetadataResponse {
	partitions := mmr.leaders[topic]
	if partitions == nil {
		partitions = make(map[int32]int32)
		mmr.leaders[topic] = partitions
	}
	partitions[partition] = brokerID
	return mmr
}

func (mmr *MockMetadataResponse) SetBroker(addr string, brokerID int32) *MockMetadataResponse {
	mmr.brokers[addr] = brokerID
	return mmr
}

func (mmr *MockMetadataResponse) For(reqBody versionedDecoder) encoder {
	metadataRequest := reqBody.(*MetadataRequest)
	metadataResponse := &MetadataResponse{}
	for addr, brokerID := range mmr.brokers {
		metadataResponse.AddBroker(addr, brokerID)
	}
	if len(metadataRequest.Topics) == 0 {
		for topic, partitions := range mmr.leaders {
			for partition, brokerID := range partitions {
				metadataResponse.AddTopicPartition(topic, partition, brokerID, nil, nil, ErrNoError)
			}
		}
		return metadataResponse
	}
	for _, topic := range metadataRequest.Topics {
		for partition, brokerID := range mmr.leaders[topic] {
			metadataResponse.AddTopicPartition(topic, partition, brokerID, nil, nil, ErrNoError)
		}
	}
	return metadataResponse
}

// MockOffsetResponse is an `OffsetResponse` builder.
type MockOffsetResponse struct {
	offsets map[string]map[int32]map[int64]int64
	t       TestReporter
}

func NewMockOffsetResponse(t TestReporter) *MockOffsetResponse {
	return &MockOffsetResponse{
		offsets: make(map[string]map[int32]map[int64]int64),
		t:       t,
	}
}

func (mor *MockOffsetResponse) SetOffset(topic string, partition int32, time, offset int64) *MockOffsetResponse {
	partitions := mor.offsets[topic]
	if partitions == nil {
		partitions = make(map[int32]map[int64]int64)
		mor.offsets[topic] = partitions
	}
	times := partitions[partition]
	if times == nil {
		times = make(map[int64]int64)
		partitions[partition] = times
	}
	times[time] = offset
	return mor
}

func (mor *MockOffsetResponse) For(reqBody versionedDecoder) encoder {
	offsetRequest := reqBody.(*OffsetRequest)
	offsetResponse := &OffsetResponse{}
	for topic, partitions := range offsetRequest.blocks {
		for partition, block := range partitions {
			offset := mor.getOffset(topic, partition, block.time)
			offsetResponse.AddTopicPartition(topic, partition, offset)
		}
	}
	return offsetResponse
}

func (mor *MockOffsetResponse) getOffset(topic string, partition int32, time int64) int64 {
	partitions := mor.offsets[topic]
	if partitions == nil {
		mor.t.Errorf("missing topic: %s", topic)
	}
	times := partitions[partition]
	if times == nil {
		mor.t.Errorf("missing partition: %d", partition)
	}
	offset, ok := times[time]
	if !ok {
		mor.t.Errorf("missing time: %d", time)
	}
	return offset
}

// MockFetchResponse is a `FetchResponse` builder.
type MockFetchResponse struct {
	messages       map[string]map[int32]map[int64]Encoder
	highWaterMarks map[string]map[int32]int64
	t              TestReporter
	batchSize      int
	version        int16
}

func NewMockFetchResponse(t TestReporter, batchSize int) *MockFetchResponse {
	return &MockFetchResponse{
		messages:       make(map[string]map[int32]map[int64]Encoder),
		highWaterMarks: make(map[string]map[int32]int64),
		t:              t,
		batchSize:      batchSize,
	}
}

func (mfr *MockFetchResponse) SetVersion(version int16) *MockFetchResponse {
	mfr.version = version
	return mfr
}

func (mfr *MockFetchResponse) SetMessage(topic string, partition int32, offset int64, msg Encoder) *MockFetchResponse {
	partitions := mfr.messages[topic]
	if partitions == nil {
		partitions = make(map[int32]map[int64]Encoder)
		mfr.messages[topic] = partitions
	}
	messages := partitions[partition]
	if messages == nil {
		messages = make(map[int64]Encoder)
		partitions[partition] = messages
	}
	messages[offset] = msg
	return mfr
}

func (mfr *MockFetchResponse) SetHighWaterMark(topic string, partition int32, offset int64) *MockFetchResponse {
	partitions := mfr.highWaterMarks[topic]
	if partitions == nil {
		partitions = make(map[int32]int64)
		mfr.highWaterMarks[topic] = partitions
	}
	partitions[partition] = offset
	return mfr
}

func (mfr *MockFetchResponse) For(reqBody versionedDecoder) encoder {
	fetchRequest := reqBody.(*FetchRequest)
	res := &FetchResponse{
		Version: mfr.version,
	}
	for topic, partitions := range fetchRequest.blocks {
		for partition, block := range partitions {
			initialOffset := block.fetchOffset
			offset := initialOffset
			maxOffset := initialOffset + int64(mfr.getMessageCount(topic, partition))
			for i := 0; i < mfr.batchSize && offset < maxOffset; {
				msg := mfr.getMessage(topic, partition, offset)
				if msg != nil {
					res.AddMessage(topic, partition, nil, msg, offset)
					i++
				}
				offset++
			}
			fb := res.GetBlock(topic, partition)
			if fb == nil {
				res.AddError(topic, partition, ErrNoError)
				fb = res.GetBlock(topic, partition)
			}
			fb.HighWaterMarkOffset = mfr.getHighWaterMark(topic, partition)
		}
	}
	return res
}

func (mfr *MockFetchResponse) getMessage(topic string, partition int32, offset int64) Encoder {
	partitions := mfr.messages[topic]
	if partitions == nil {
		return nil
	}
	messages := partitions[partition]
	if messages == nil {
		return nil
	}
	return messages[offset]
}

func (mfr *MockFetchResponse) getMessageCount(topic string, partition int32) int {
	partitions := mfr.messages[topic]
	if partitions == nil {
		return 0
	}
	messages := partitions[partition]
	if messages == nil {
		return 0
	}
	return len(messages)
}

func (mfr *MockFetchResponse) getHighWaterMark(topic string, partition int32) int64 {
	partitions := mfr.highWaterMarks[topic]
	if partitions == nil {
		return 0
	}
	return partitions[partition]
}

// MockConsumerMetadataResponse is a `ConsumerMetadataResponse` builder.
type MockConsumerMetadataResponse struct {
	coordinators map[string]interface{}
	t            TestReporter
}

func NewMockConsumerMetadataResponse(t TestReporter) *MockConsumerMetadataResponse {
	return &MockConsumerMetadataResponse{
		coordinators: make(map[string]interface{}),
		t:            t,
	}
}

func (mr *MockConsumerMetadataResponse) SetCoordinator(group string, broker *MockBroker) *MockConsumerMetadataResponse {
	mr.coordinators[group] = broker
	return mr
}

func (mr *MockConsumerMetadataResponse) SetError(group string, kerror KError) *MockConsumerMetadataResponse {
	mr.coordinators[group] = kerror
	return mr
}

func (mr *MockConsumerMetadataResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*ConsumerMetadataRequest)
	group := req.ConsumerGroup
	res := &ConsumerMetadataResponse{}
	v := mr.coordinators[group]
	switch v := v.(type) {
	case *MockBroker:
		res.Coordinator = &Broker{id: v.BrokerID(), addr: v.Addr()}
	case KError:
		res.Err = v
	}
	return res
}

// MockOffsetCommitResponse is a `OffsetCommitResponse` builder.
type MockOffsetCommitResponse struct {
	errors map[string]map[string]map[int32]KError
	t      TestReporter
}

func NewMockOffsetCommitResponse(t TestReporter) *MockOffsetCommitResponse {
	return &MockOffsetCommitResponse{t: t}
}

func (mr *MockOffsetCommitResponse) SetError(group, topic string, partition int32, kerror KError) *MockOffsetCommitResponse {
	if mr.errors == nil {
		mr.errors = make(map[string]map[string]map[int32]KError)
	}
	topics := mr.errors[group]
	if topics == nil {
		topics = make(map[string]map[int32]KError)
		mr.errors[group] = topics
	}
	partitions := topics[topic]
	if partitions == nil {
		partitions = make(map[int32]KError)
		topics[topic] = partitions
	}
	partitions[partition] = kerror
	return mr
}

func (mr *MockOffsetCommitResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*OffsetCommitRequest)
	group := req.ConsumerGroup
	res := &OffsetCommitResponse{}
	for topic, partitions := range req.blocks {
		for partition := range partitions {
			res.AddError(topic, partition, mr.getError(group, topic, partition))
		}
	}
	return res
}

func (mr *MockOffsetCommitResponse) getError(group, topic string, partition int32) KError {
	topics := mr.errors[group]
	if topics == nil {
		return ErrNoError
	}
	partitions := topics[topic]
	if partitions == nil {
		return ErrNoError
	}
	kerror, ok := partitions[partition]
	if !ok {
		return ErrNoError
	}
	return kerror
}

// MockProduceResponse is a `ProduceResponse` builder.
type MockProduceResponse struct {
	errors map[string]map[int32]KError
	t      TestReporter
}

func NewMockProduceResponse(t TestReporter) *MockProduceResponse {
	return &MockProduceResponse{t: t}
}

func (mr *MockProduceResponse) SetError(topic string, partition int32, kerror KError) *MockProduceResponse {
	if mr.errors == nil {
		mr.errors = make(map[string]map[int32]KError)
	}
	partitions := mr.errors[topic]
	if partitions == nil {
		partitions = make(map[int32]KError)
		mr.errors[topic] = partitions
	}
	partitions[partition] = kerror
	return mr
}

func (mr *MockProduceResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*ProduceRequest)
	res := &ProduceResponse{}
	for topic, partitions := range req.msgSets {
		for partition := range partitions {
			res.AddTopicPartition(topic, partition, mr.getError(topic, partition))
		}
	}
	return res
}

func (mr *MockProduceResponse) getError(topic string, partition int32) KError {
	partitions := mr.errors[topic]
	if partitions == nil {
		return ErrNoError
	}
	kerror, ok := partitions[partition]
	if !ok {
		return ErrNoError
	}
	return kerror
}

// MockOffsetFetchResponse is a `OffsetFetchResponse` builder.
type MockOffsetFetchResponse struct {
	offsets map[string]map[string]map[int32]*OffsetFetchResponseBlock
	t       TestReporter
}

func NewMockOffsetFetchResponse(t TestReporter) *MockOffsetFetchResponse {
	return &MockOffsetFetchResponse{t: t}
}

func (mr *MockOffsetFetchResponse) SetOffset(group, topic string, partition int32, offset int64, metadata string, kerror KError) *MockOffsetFetchResponse {
	if mr.offsets == nil {
		mr.offsets = make(map[string]map[string]map[int32]*OffsetFetchResponseBlock)
	}
	topics := mr.offsets[group]
	if topics == nil {
		topics = make(map[string]map[int32]*OffsetFetchResponseBlock)
		mr.offsets[group] = topics
	}
	partitions := topics[topic]
	if partitions == nil {
		partitions = make(map[int32]*OffsetFetchResponseBlock)
		topics[topic] = partitions
	}
	partitions[partition] = &OffsetFetchResponseBlock{offset, metadata, kerror}
	return mr
}

func (mr *MockOffsetFetchResponse) For(reqBody versionedDecoder) encoder {
	req := reqBody.(*OffsetFetchRequest)
	group := req.ConsumerGroup
	res := &OffsetFetchResponse{}
	for topic, partitions := range mr.offsets[group] {
		for partition, block := range partitions {
			res.AddBlock(topic, partition, block)
		}
	}
	return res
}
<|endoftext|>


<commit_before>package models

import (
	"fmt"
	"net"
)

// Config ...
type Config struct {
	Provider string `json:"provider"`
	CIMode   bool   `json:"ci-mode"`

	// required for docker-machine
	MountType      string `json:"mount-type"`
	NetfsMountOpts string `json:"netfs-mount-opts"`
	CPUs           int    `json:"cpus"`
	RAM            int    `json:"ram"`
	Disk           int    `json:"disk"`

	// ip address spaces
	ExternalNetworkSpace      string `json:"external-network-space"`
	DockerMachineNetworkSpace string `json:"docker-machine-network-space"`
	NativeNetworkSpace        string `json:"native-network-space"`

	LockPort int `json:"lock-port"`
}

// Save persists the Config to the database
func (c *Config) Save() error {
	// make sure the information in is valid
	c.makeValid()

	// Since there is only ever a single Config value, we'll use the registry
	if err := put("registry", "Config", c); err != nil {
		return fmt.Errorf("failed to save Config: %s", err.Error())
	}

	return nil
}

// set reasonable default values for all necessary values
func (c *Config) makeValid() {
	if c.Provider != "native" && c.Provider != "docker-machine" {
		c.Provider = "docker-machine"
	}

	if c.MountType != "native" && c.MountType != "netfs" {
		c.MountType = "native"
	}

	if c.CPUs == 0 {
		c.CPUs = 1
	}

	if c.RAM == 0 {
		c.RAM = 1
	}

	if c.Disk < 102400 {
		c.Disk = 102400
	}

	if _, _, err := net.ParseCIDR(c.ExternalNetworkSpace); c.ExternalNetworkSpace == "" || err != nil {
		c.ExternalNetworkSpace = "192.168.99.50/24"
	}

	if _, _, err := net.ParseCIDR(c.DockerMachineNetworkSpace); c.DockerMachineNetworkSpace == "" || err != nil {
		c.DockerMachineNetworkSpace = "172.19.0.1/16"
	}

	if _, _, err := net.ParseCIDR(c.NativeNetworkSpace); c.NativeNetworkSpace == "" || err != nil {
		c.NativeNetworkSpace = "172.18.0.1/16"
	}

	if c.LockPort == 0 {
		c.LockPort = 12345
	}

}

// Delete deletes the Config record from the database
func (c *Config) Delete() error {

	// Since there is only ever a single Config value, we'll use the registry
	if err := destroy("registry", "Config"); err != nil {
		return fmt.Errorf("failed to delete Config: %s", err.Error())
	}

	// clear the current entry
	c = nil

	return nil
}

// LoadConfig loads the Config entry
func LoadConfig() (*Config, error) {
	c := &Config{}
	c.makeValid()
	if err := get("registry", "Config", &c); err != nil {
		return c, fmt.Errorf("failed to load Config: %s", err.Error())
	}

	return c, nil
}
<commit_msg>set a default for netfs on windows fixes #422<commit_after>package models

import (
	"fmt"
	"net"
	"runtime"
)

// Config ...
type Config struct {
	Provider string `json:"provider"`
	CIMode   bool   `json:"ci-mode"`

	// required for docker-machine
	MountType      string `json:"mount-type"`
	NetfsMountOpts string `json:"netfs-mount-opts"`
	CPUs           int    `json:"cpus"`
	RAM            int    `json:"ram"`
	Disk           int    `json:"disk"`

	// ip address spaces
	ExternalNetworkSpace      string `json:"external-network-space"`
	DockerMachineNetworkSpace string `json:"docker-machine-network-space"`
	NativeNetworkSpace        string `json:"native-network-space"`

	LockPort int `json:"lock-port"`
}

// Save persists the Config to the database
func (c *Config) Save() error {
	// make sure the information in is valid
	c.makeValid()

	// Since there is only ever a single Config value, we'll use the registry
	if err := put("registry", "Config", c); err != nil {
		return fmt.Errorf("failed to save Config: %s", err.Error())
	}

	return nil
}

// set reasonable default values for all necessary values
func (c *Config) makeValid() {
	if c.Provider != "native" && c.Provider != "docker-machine" {
		c.Provider = "docker-machine"
	}

	if c.MountType != "native" && c.MountType != "netfs" {
		c.MountType = "native"
	}

	if c.CPUs == 0 {
		c.CPUs = 1
	}

	if c.RAM == 0 {
		c.RAM = 1
	}

	if c.Disk < 102400 {
		c.Disk = 102400
	}

	if c.NetfsMountOpts == "" && runtime.GOOS == "windows" {
		c.NetfsMountOpts = "mfsymlinks"
	}

	if _, _, err := net.ParseCIDR(c.ExternalNetworkSpace); c.ExternalNetworkSpace == "" || err != nil {
		c.ExternalNetworkSpace = "192.168.99.50/24"
	}

	if _, _, err := net.ParseCIDR(c.DockerMachineNetworkSpace); c.DockerMachineNetworkSpace == "" || err != nil {
		c.DockerMachineNetworkSpace = "172.19.0.1/16"
	}

	if _, _, err := net.ParseCIDR(c.NativeNetworkSpace); c.NativeNetworkSpace == "" || err != nil {
		c.NativeNetworkSpace = "172.18.0.1/16"
	}

	if c.LockPort == 0 {
		c.LockPort = 12345
	}

}

// Delete deletes the Config record from the database
func (c *Config) Delete() error {

	// Since there is only ever a single Config value, we'll use the registry
	if err := destroy("registry", "Config"); err != nil {
		return fmt.Errorf("failed to delete Config: %s", err.Error())
	}

	// clear the current entry
	c = nil

	return nil
}

// LoadConfig loads the Config entry
func LoadConfig() (*Config, error) {
	c := &Config{}
	c.makeValid()
	if err := get("registry", "Config", &c); err != nil {
		return c, fmt.Errorf("failed to load Config: %s", err.Error())
	}

	return c, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"time"

	log "github.com/coreos/fleet/Godeps/_workspace/src/github.com/golang/glog"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/registry"
)

const (
	// time between triggering reconciliation routine
	reconcileInterval = 5 * time.Second

	taskTypeLoadJob   = "LoadJob"
	taskTypeUnloadJob = "UnloadJob"
	taskTypeStartJob  = "StartJob"
	taskTypeStopJob   = "StopJob"

	taskReasonScheduledButNotRunnable    = "job scheduled locally but unable to run"
	taskReasonScheduledButUnloaded       = "job scheduled here but not loaded"
	taskReasonLoadedButNotScheduled      = "job loaded but not scheduled here"
	taskReasonLoadedDesiredStateLaunched = "job currently loaded but desired state is launched"
	taskReasonLaunchedDesiredStateLoaded = "job currently launched but desired state is loaded"
	taskReasonPurgingAgent               = "purging agent"
)

type task struct {
	Type   string
	Job    *job.Job
	Reason string
}

func (t *task) String() string {
	var jName string
	if t.Job != nil {
		jName = t.Job.Name
	}
	return fmt.Sprintf("{Type: %s, Job: %s, Reason: %q}", t.Type, jName, t.Reason)
}

func NewReconciler(reg registry.Registry, rStream registry.EventStream) (*AgentReconciler, error) {
	ar := AgentReconciler{reg, rStream}
	return &ar, nil
}

type AgentReconciler struct {
	reg     registry.Registry
	rStream registry.EventStream
}

// Run periodically attempts to reconcile the provided Agent until the stop
// channel is closed. Run will also reconcile in reaction to calls to Trigger.
// While a reconciliation is being attempted, calls to Trigger are ignored.
func (ar *AgentReconciler) Run(a *Agent, stop chan bool) {
	ticker := time.Tick(reconcileInterval)

	reconcile := func() {
		start := time.Now()
		ar.Reconcile(a)
		elapsed := time.Now().Sub(start)

		msg := fmt.Sprintf("AgentReconciler completed reconciliation in %s", elapsed)
		if elapsed > reconcileInterval {
			log.Warning(msg)
		} else {
			log.V(1).Info(msg)
		}
	}

	trigger := make(chan struct{})
	go func() {
		abort := make(chan struct{})
		select {
		case <-stop:
			close(abort)
		case <-ar.rStream.Next(abort):
			trigger <- struct{}{}
		}
	}()

	for {
		select {
		case <-stop:
			log.V(1).Info("AgentReconciler exiting due to stop signal")
			return
		case <-ticker:
			reconcile()
		case <-trigger:
			reconcile()
		}
	}
}

// Reconcile drives the local Agent's state towards the desired state
// stored in the Registry.
func (ar *AgentReconciler) Reconcile(a *Agent) {
	ms := a.Machine.State()

	units, err := ar.reg.Units()
	if err != nil {
		log.Errorf("Failed fetching Units from Registry: %v", err)
		return
	}

	sUnits, err := ar.reg.Schedule()
	if err != nil {
		log.Errorf("Failed fetching schedule from Registry: %v", err)
		return
	}

	dAgentState, err := ar.desiredAgentState(units, sUnits, &ms)
	if err != nil {
		log.Errorf("Unable to determine agent's desired state: %v", err)
		return
	}

	cAgentState, err := ar.currentAgentState(a)
	if err != nil {
		log.Errorf("Unable to determine agent's current state: %v", err)
		return
	}

	for t := range ar.calculateTasksForJobs(dAgentState, cAgentState) {
		err := ar.doTask(a, t)
		if err != nil {
			log.Errorf("Failed resolving task, halting reconciliation: task=%s err=%q", t, err)
			return
		}
	}
}

// Purge attempts to unload all Jobs that have been loaded locally
func (ar *AgentReconciler) Purge(a *Agent) {
	cAgentState, err := ar.currentAgentState(a)
	if err != nil {
		log.Errorf("Unable to determine agent's current state: %v", err)
		return
	}

	for _, cJob := range cAgentState.Jobs {
		t := task{
			Type:   taskTypeUnloadJob,
			Job:    cJob,
			Reason: taskReasonPurgingAgent,
		}

		err := ar.doTask(a, &t)
		if err != nil {
			log.Errorf("Failed resolving task: task=%s err=%q", t, err)
		}
	}
}

// doTask takes action on an Agent based on the contents of a *task
func (ar *AgentReconciler) doTask(a *Agent, t *task) (err error) {
	switch t.Type {
	case taskTypeLoadJob:
		err = a.loadJob(t.Job)
	case taskTypeUnloadJob:
		a.unloadJob(t.Job.Name)
	case taskTypeStartJob:
		a.startJob(t.Job.Name)
	case taskTypeStopJob:
		a.stopJob(t.Job.Name)
	default:
		err = fmt.Errorf("unrecognized task type %q", t.Type)
	}

	if err == nil {
		log.Infof("AgentReconciler completed task: %s", t)
	}

	return
}

// desiredAgentState builds an *AgentState object that represents what an
// Agent identified by the provided machine ID should currently be doing.
func (ar *AgentReconciler) desiredAgentState(units []job.Unit, sUnits []job.ScheduledUnit, ms *machine.MachineState) (*AgentState, error) {
	as := AgentState{
		MState: ms,
		Jobs:   make(map[string]*job.Job),
	}

	sUnitMap := make(map[string]*job.ScheduledUnit)
	for _, sUnit := range sUnits {
		sUnit := sUnit
		sUnitMap[sUnit.Name] = &sUnit
	}

	for _, u := range units {
		sUnit, ok := sUnitMap[u.Name]
		if !ok || sUnit.TargetMachineID == "" || sUnit.TargetMachineID != ms.ID {
			continue
		}

		as.Jobs[u.Name] = &job.Job{
			Name:            u.Name,
			Unit:            u.Unit,
			TargetState:     u.TargetState,
			TargetMachineID: sUnit.TargetMachineID,
			State:           sUnit.State,
		}
	}

	return &as, nil
}

// currentAgentState builds an *AgentState object that represents what an
// Agent is currently doing.
func (ar *AgentReconciler) currentAgentState(a *Agent) (*AgentState, error) {
	jobs, err := a.jobs()
	if err != nil {
		return nil, err
	}

	ms := a.Machine.State()
	as := AgentState{
		MState: &ms,
		Jobs:   jobs,
	}

	return &as, nil
}

// calculateTasksForJobs compares the desired and current state of an Agent.
// The generateed tasks represent what should be done to make the desired
// state match the current state.
func (ar *AgentReconciler) calculateTasksForJobs(dState, cState *AgentState) <-chan *task {
	taskchan := make(chan *task)
	go func() {
		jobs := pkg.NewUnsafeSet()
		for cName := range cState.Jobs {
			jobs.Add(cName)
		}

		for dName := range dState.Jobs {
			jobs.Add(dName)
		}

		for _, name := range jobs.Values() {
			ar.calculateTasksForJob(dState, cState, name, taskchan)
		}

		close(taskchan)
	}()

	return taskchan
}

func (ar *AgentReconciler) calculateTasksForJob(dState, cState *AgentState, jName string, taskchan chan *task) {
	var dJob, cJob *job.Job
	if dState != nil {
		dJob = dState.Jobs[jName]
	}
	if cState != nil {
		cJob = cState.Jobs[jName]
	}

	if dJob == nil && cJob == nil {
		log.Errorf("Desired state and current state of Job(%s) nil, not sure what to do", jName)
		return
	}

	if dJob == nil || dJob.TargetState == job.JobStateInactive {
		taskchan <- &task{
			Type:   taskTypeUnloadJob,
			Job:    cJob,
			Reason: taskReasonLoadedButNotScheduled,
		}

		delete(cState.Jobs, jName)
		return
	}

	if cJob == nil {
		taskchan <- &task{
			Type:   taskTypeLoadJob,
			Job:    dJob,
			Reason: taskReasonScheduledButUnloaded,
		}

		return
	}

	if cJob.State == nil {
		log.Errorf("Current state of Job(%s) unknown, unable to reconcile", jName)
		return
	}

	if dJob.State == nil {
		log.Errorf("Desired state of Job(%s) unknown, unable to reconcile", jName)
		return
	}

	if *cJob.State == dJob.TargetState {
		log.V(1).Infof("Desired state %q matches current state of Job(%s), nothing to do", *cJob.State, jName)
		return
	}

	if *cJob.State == job.JobStateInactive {
		taskchan <- &task{
			Type:   taskTypeLoadJob,
			Job:    dJob,
			Reason: taskReasonScheduledButUnloaded,
		}
	}

	if (*cJob.State == job.JobStateInactive || *cJob.State == job.JobStateLoaded) && dJob.TargetState == job.JobStateLaunched {
		taskchan <- &task{
			Type:   taskTypeStartJob,
			Job:    cJob,
			Reason: taskReasonLoadedDesiredStateLaunched,
		}
		return
	}

	if *cJob.State == job.JobStateLaunched && dJob.TargetState == job.JobStateLoaded {
		taskchan <- &task{
			Type:   taskTypeStopJob,
			Job:    cJob,
			Reason: taskReasonLaunchedDesiredStateLoaded,
		}
		return
	}

	log.Errorf("Unable to determine how to reconcile Job(%s): desiredState=%#v currentState=%#V", jName, dJob, cJob)
}
<commit_msg>agent/reconcile: remove State from desired state<commit_after>package agent

import (
	"fmt"
	"time"

	log "github.com/coreos/fleet/Godeps/_workspace/src/github.com/golang/glog"

	"github.com/coreos/fleet/job"
	"github.com/coreos/fleet/machine"
	"github.com/coreos/fleet/pkg"
	"github.com/coreos/fleet/registry"
)

const (
	// time between triggering reconciliation routine
	reconcileInterval = 5 * time.Second

	taskTypeLoadJob   = "LoadJob"
	taskTypeUnloadJob = "UnloadJob"
	taskTypeStartJob  = "StartJob"
	taskTypeStopJob   = "StopJob"

	taskReasonScheduledButNotRunnable    = "job scheduled locally but unable to run"
	taskReasonScheduledButUnloaded       = "job scheduled here but not loaded"
	taskReasonLoadedButNotScheduled      = "job loaded but not scheduled here"
	taskReasonLoadedDesiredStateLaunched = "job currently loaded but desired state is launched"
	taskReasonLaunchedDesiredStateLoaded = "job currently launched but desired state is loaded"
	taskReasonPurgingAgent               = "purging agent"
)

type task struct {
	Type   string
	Job    *job.Job
	Reason string
}

func (t *task) String() string {
	var jName string
	if t.Job != nil {
		jName = t.Job.Name
	}
	return fmt.Sprintf("{Type: %s, Job: %s, Reason: %q}", t.Type, jName, t.Reason)
}

func NewReconciler(reg registry.Registry, rStream registry.EventStream) (*AgentReconciler, error) {
	ar := AgentReconciler{reg, rStream}
	return &ar, nil
}

type AgentReconciler struct {
	reg     registry.Registry
	rStream registry.EventStream
}

// Run periodically attempts to reconcile the provided Agent until the stop
// channel is closed. Run will also reconcile in reaction to calls to Trigger.
// While a reconciliation is being attempted, calls to Trigger are ignored.
func (ar *AgentReconciler) Run(a *Agent, stop chan bool) {
	ticker := time.Tick(reconcileInterval)

	reconcile := func() {
		start := time.Now()
		ar.Reconcile(a)
		elapsed := time.Now().Sub(start)

		msg := fmt.Sprintf("AgentReconciler completed reconciliation in %s", elapsed)
		if elapsed > reconcileInterval {
			log.Warning(msg)
		} else {
			log.V(1).Info(msg)
		}
	}

	trigger := make(chan struct{})
	go func() {
		abort := make(chan struct{})
		select {
		case <-stop:
			close(abort)
		case <-ar.rStream.Next(abort):
			trigger <- struct{}{}
		}
	}()

	for {
		select {
		case <-stop:
			log.V(1).Info("AgentReconciler exiting due to stop signal")
			return
		case <-ticker:
			reconcile()
		case <-trigger:
			reconcile()
		}
	}
}

// Reconcile drives the local Agent's state towards the desired state
// stored in the Registry.
func (ar *AgentReconciler) Reconcile(a *Agent) {
	ms := a.Machine.State()

	units, err := ar.reg.Units()
	if err != nil {
		log.Errorf("Failed fetching Units from Registry: %v", err)
		return
	}

	sUnits, err := ar.reg.Schedule()
	if err != nil {
		log.Errorf("Failed fetching schedule from Registry: %v", err)
		return
	}

	dAgentState, err := ar.desiredAgentState(units, sUnits, &ms)
	if err != nil {
		log.Errorf("Unable to determine agent's desired state: %v", err)
		return
	}

	cAgentState, err := ar.currentAgentState(a)
	if err != nil {
		log.Errorf("Unable to determine agent's current state: %v", err)
		return
	}

	for t := range ar.calculateTasksForJobs(dAgentState, cAgentState) {
		err := ar.doTask(a, t)
		if err != nil {
			log.Errorf("Failed resolving task, halting reconciliation: task=%s err=%q", t, err)
			return
		}
	}
}

// Purge attempts to unload all Jobs that have been loaded locally
func (ar *AgentReconciler) Purge(a *Agent) {
	cAgentState, err := ar.currentAgentState(a)
	if err != nil {
		log.Errorf("Unable to determine agent's current state: %v", err)
		return
	}

	for _, cJob := range cAgentState.Jobs {
		t := task{
			Type:   taskTypeUnloadJob,
			Job:    cJob,
			Reason: taskReasonPurgingAgent,
		}

		err := ar.doTask(a, &t)
		if err != nil {
			log.Errorf("Failed resolving task: task=%s err=%q", t, err)
		}
	}
}

// doTask takes action on an Agent based on the contents of a *task
func (ar *AgentReconciler) doTask(a *Agent, t *task) (err error) {
	switch t.Type {
	case taskTypeLoadJob:
		err = a.loadJob(t.Job)
	case taskTypeUnloadJob:
		a.unloadJob(t.Job.Name)
	case taskTypeStartJob:
		a.startJob(t.Job.Name)
	case taskTypeStopJob:
		a.stopJob(t.Job.Name)
	default:
		err = fmt.Errorf("unrecognized task type %q", t.Type)
	}

	if err == nil {
		log.Infof("AgentReconciler completed task: %s", t)
	}

	return
}

// desiredAgentState builds an *AgentState object that represents what an
// Agent identified by the provided machine ID should currently be doing.
func (ar *AgentReconciler) desiredAgentState(units []job.Unit, sUnits []job.ScheduledUnit, ms *machine.MachineState) (*AgentState, error) {
	as := AgentState{
		MState: ms,
		Jobs:   make(map[string]*job.Job),
	}

	sUnitMap := make(map[string]*job.ScheduledUnit)
	for _, sUnit := range sUnits {
		sUnit := sUnit
		sUnitMap[sUnit.Name] = &sUnit
	}

	for _, u := range units {
		sUnit, ok := sUnitMap[u.Name]
		if !ok || sUnit.TargetMachineID == "" || sUnit.TargetMachineID != ms.ID {
			continue
		}

		as.Jobs[u.Name] = &job.Job{
			Name:            u.Name,
			Unit:            u.Unit,
			TargetState:     u.TargetState,
			TargetMachineID: sUnit.TargetMachineID,
		}
	}

	return &as, nil
}

// currentAgentState builds an *AgentState object that represents what an
// Agent is currently doing.
func (ar *AgentReconciler) currentAgentState(a *Agent) (*AgentState, error) {
	jobs, err := a.jobs()
	if err != nil {
		return nil, err
	}

	ms := a.Machine.State()
	as := AgentState{
		MState: &ms,
		Jobs:   jobs,
	}

	return &as, nil
}

// calculateTasksForJobs compares the desired and current state of an Agent.
// The generateed tasks represent what should be done to make the desired
// state match the current state.
func (ar *AgentReconciler) calculateTasksForJobs(dState, cState *AgentState) <-chan *task {
	taskchan := make(chan *task)
	go func() {
		jobs := pkg.NewUnsafeSet()
		for cName := range cState.Jobs {
			jobs.Add(cName)
		}

		for dName := range dState.Jobs {
			jobs.Add(dName)
		}

		for _, name := range jobs.Values() {
			ar.calculateTasksForJob(dState, cState, name, taskchan)
		}

		close(taskchan)
	}()

	return taskchan
}

func (ar *AgentReconciler) calculateTasksForJob(dState, cState *AgentState, jName string, taskchan chan *task) {
	var dJob, cJob *job.Job
	if dState != nil {
		dJob = dState.Jobs[jName]
	}
	if cState != nil {
		cJob = cState.Jobs[jName]
	}

	if dJob == nil && cJob == nil {
		log.Errorf("Desired state and current state of Job(%s) nil, not sure what to do", jName)
		return
	}

	if dJob == nil || dJob.TargetState == job.JobStateInactive {
		taskchan <- &task{
			Type:   taskTypeUnloadJob,
			Job:    cJob,
			Reason: taskReasonLoadedButNotScheduled,
		}

		delete(cState.Jobs, jName)
		return
	}

	if cJob == nil {
		taskchan <- &task{
			Type:   taskTypeLoadJob,
			Job:    dJob,
			Reason: taskReasonScheduledButUnloaded,
		}

		return
	}

	if cJob.State == nil {
		log.Errorf("Current state of Job(%s) unknown, unable to reconcile", jName)
		return
	}

	if *cJob.State == dJob.TargetState {
		log.V(1).Infof("Desired state %q matches current state of Job(%s), nothing to do", *cJob.State, jName)
		return
	}

	if *cJob.State == job.JobStateInactive {
		taskchan <- &task{
			Type:   taskTypeLoadJob,
			Job:    dJob,
			Reason: taskReasonScheduledButUnloaded,
		}
	}

	if (*cJob.State == job.JobStateInactive || *cJob.State == job.JobStateLoaded) && dJob.TargetState == job.JobStateLaunched {
		taskchan <- &task{
			Type:   taskTypeStartJob,
			Job:    cJob,
			Reason: taskReasonLoadedDesiredStateLaunched,
		}
		return
	}

	if *cJob.State == job.JobStateLaunched && dJob.TargetState == job.JobStateLoaded {
		taskchan <- &task{
			Type:   taskTypeStopJob,
			Job:    cJob,
			Reason: taskReasonLaunchedDesiredStateLoaded,
		}
		return
	}

	log.Errorf("Unable to determine how to reconcile Job(%s): desiredState=%#v currentState=%#V", jName, dJob, cJob)
}
<|endoftext|>


<commit_before>// Copyright 2015 monsterqueue authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongodb

import (
	"errors"
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/tsuru/monsterqueue"
	"github.com/tsuru/monsterqueue/log"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type queueMongoDB struct {
	config   *QueueConfig
	session  *mgo.Session
	tasks    map[string]monsterqueue.Task
	tasksMut sync.RWMutex
	done     chan bool
	wg       sync.WaitGroup
}

type QueueConfig struct {
	Url              string // MongoDB connection url
	Database         string // MongoDB database name
	CollectionPrefix string // Prefix for all collections created in MongoDB
}

// Creates a new queue. The QueueConfig parameter will tell us how to connect
// to mongodb. This command will fail if the MongoDB server is not available.
//
// Tasks registered in this queue instance will run when `ProcessLoop` is
// called in this *same* instance.
func NewQueue(conf QueueConfig) (monsterqueue.Queue, error) {
	q := &queueMongoDB{
		config: &conf,
		tasks:  make(map[string]monsterqueue.Task),
		done:   make(chan bool),
	}
	var err error
	if conf.Url == "" {
		return nil, errors.New("setting QueueConfig.Url is required")
	}
	dialInfo, err := mgo.ParseURL(conf.Url)
	if err != nil {
		return nil, err
	}
	dialInfo.FailFast = true
	q.session, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		return nil, err
	}
	q.session.SetSyncTimeout(10 * time.Second)
	q.session.SetSocketTimeout(1 * time.Minute)
	db := q.session.DB(conf.Database)
	if db.Name == "test" {
		q.session.Close()
		return nil, errors.New("database name should be set in QueueConfig.Url or QueueConfig.Database")
	}
	return q, err
}

func (q *queueMongoDB) tasksColl() *mgo.Collection {
	s := q.session.Copy()
	name := "queue_tasks"
	if q.config.CollectionPrefix != "" {
		name = fmt.Sprintf("%s_%s", q.config.CollectionPrefix, name)
	}
	return s.DB(q.config.Database).C(name)
}

func (q *queueMongoDB) RegisterTask(task monsterqueue.Task) error {
	q.tasksMut.Lock()
	defer q.tasksMut.Unlock()
	if _, isRegistered := q.tasks[task.Name()]; isRegistered {
		return errors.New("task already registered")
	}
	q.tasks[task.Name()] = task
	return nil
}

func (q *queueMongoDB) Enqueue(taskName string, params monsterqueue.JobParams) (monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	j := q.initialJob(taskName, params)
	err := coll.Insert(j)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

func (q *queueMongoDB) getDoneJob(jobId bson.ObjectId) (*jobMongoDB, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var resultJob jobMongoDB
	err := coll.Find(bson.M{"_id": jobId, "resultmessage.done": true}).One(&resultJob)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, nil
		}
		return nil, err
	}
	return &resultJob, nil
}

func (q *queueMongoDB) EnqueueWait(taskName string, params monsterqueue.JobParams, timeout time.Duration) (monsterqueue.Job, error) {
	j := q.initialJob(taskName, params)
	j.Waited = true
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	err := coll.Insert(j)
	if err != nil {
		return nil, err
	}
	result := make(chan *jobMongoDB)
	quit := make(chan bool)
	go func() {
		for range time.Tick(200 * time.Millisecond) {
			select {
			case <-quit:
				return
			default:
			}
			job, err := q.getDoneJob(j.Id)
			if err != nil {
				log.Errorf("error trying to get job %s: %s", job.Id, err.Error())
			}
			if job != nil {
				result <- job
				return
			}
		}
	}()
	var resultJob *jobMongoDB
	select {
	case resultJob = <-result:
	case <-time.After(timeout):
		quit <- true
	}
	close(quit)
	close(result)
	if resultJob != nil {
		return resultJob, nil
	}
	err = coll.Update(bson.M{
		"_id":    j.Id,
		"waited": true,
	}, bson.M{"$set": bson.M{"waited": false}})
	if err == mgo.ErrNotFound {
		resultJob, err = q.getDoneJob(j.Id)
	}
	if err != nil {
		return &j, err
	}
	if resultJob != nil {
		return resultJob, nil
	}
	return &j, monsterqueue.ErrQueueWaitTimeout
}

func (q *queueMongoDB) ProcessLoop() {
	for {
		q.wg.Add(1)
		err := q.waitForMessage()
		if err != nil {
			log.Debugf("error getting message from queue: %s", err.Error())
		}
		select {
		case <-time.After(1 * time.Second):
		case <-q.done:
			return
		}
	}
}

func (q *queueMongoDB) Stop() {
	q.done <- true
	q.Wait()
}

func (q *queueMongoDB) Wait() {
	q.wg.Wait()
}

func (q *queueMongoDB) ResetStorage() error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.DropCollection()
}

func (q *queueMongoDB) RetrieveJob(jobId string) (monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var job jobMongoDB
	err := coll.FindId(bson.ObjectIdHex(jobId)).One(&job)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, monsterqueue.ErrNoSuchJob
		}
		return nil, err
	}
	return &job, err
}

func (q *queueMongoDB) ListJobs() ([]monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var mongodbJobs []jobMongoDB
	err := coll.Find(nil).All(&mongodbJobs)
	if err != nil {
		return nil, err
	}
	jobs := make([]monsterqueue.Job, len(mongodbJobs))
	for i := range mongodbJobs {
		jobs[i] = &mongodbJobs[i]
	}
	return jobs, nil
}

func (q *queueMongoDB) DeleteJob(jobId string) error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.RemoveId(bson.ObjectIdHex(jobId))
}

func (q *queueMongoDB) initialJob(taskName string, params monsterqueue.JobParams) jobMongoDB {
	buf := make([]byte, monsterqueue.StackTraceLimit)
	buf = buf[:runtime.Stack(buf, false)]
	return jobMongoDB{
		Id:        bson.NewObjectId(),
		Task:      taskName,
		Params:    params,
		Timestamp: time.Now().UTC(),
		Stack:     string(buf),
		queue:     q,
	}
}

func (q *queueMongoDB) waitForMessage() error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var job jobMongoDB
	hostname, _ := os.Hostname()
	ownerData := jobOwnership{
		Name:      fmt.Sprintf("%s_%d", hostname, os.Getpid()),
		Owned:     true,
		Timestamp: time.Now().UTC(),
	}
	q.tasksMut.RLock()
	taskNames := make([]string, 0, len(q.tasks))
	for taskName := range q.tasks {
		taskNames = append(taskNames, taskName)
	}
	q.tasksMut.RUnlock()
	_, err := coll.Find(bson.M{
		"task":               bson.M{"$in": taskNames},
		"owner.owned":        false,
		"resultmessage.done": false,
	}).Sort("_id").Apply(mgo.Change{
		Update: bson.M{
			"$set": bson.M{"owner": ownerData},
		},
	}, &job)
	if err != nil {
		q.wg.Done()
		if err == mgo.ErrNotFound {
			return nil
		}
		return err
	}
	job.queue = q
	if err != nil {
		q.moveToResult(&job, nil, err)
		q.wg.Done()
		return err
	}
	q.tasksMut.RLock()
	task, _ := q.tasks[job.Task]
	q.tasksMut.RUnlock()
	if task == nil {
		err := fmt.Errorf("unregistered task name %q", job.Task)
		q.moveToResult(&job, nil, err)
		q.wg.Done()
		return err
	}
	go func() {
		defer q.wg.Done()
		task.Run(&job)
		if !job.ResultMessage.Done {
			q.moveToResult(&job, nil, monsterqueue.ErrNoJobResultSet)
		}
	}()
	return nil
}

func (q *queueMongoDB) moveToResult(job *jobMongoDB, result monsterqueue.JobResult, jobErr error) error {
	var resultMsg jobResultMessage
	resultMsg.Result = result
	resultMsg.Timestamp = time.Now().UTC()
	resultMsg.Done = true
	if jobErr != nil {
		resultMsg.Error = jobErr.Error()
	}
	job.ResultMessage = resultMsg
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.UpdateId(job.Id, bson.M{"$set": bson.M{"resultmessage": resultMsg, "owner.owned": false}})
}

func (q *queueMongoDB) publishResult(job *jobMongoDB) (bool, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	err := coll.Update(bson.M{"_id": job.Id, "waited": true}, bson.M{"$set": bson.M{"waited": false}})
	if err != nil {
		if err == mgo.ErrNotFound {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<commit_msg>mongodb: do not use the job when getDoneJob fails<commit_after>// Copyright 2015 monsterqueue authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mongodb

import (
	"errors"
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"

	"github.com/tsuru/monsterqueue"
	"github.com/tsuru/monsterqueue/log"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

type queueMongoDB struct {
	config   *QueueConfig
	session  *mgo.Session
	tasks    map[string]monsterqueue.Task
	tasksMut sync.RWMutex
	done     chan bool
	wg       sync.WaitGroup
}

type QueueConfig struct {
	Url              string // MongoDB connection url
	Database         string // MongoDB database name
	CollectionPrefix string // Prefix for all collections created in MongoDB
}

// Creates a new queue. The QueueConfig parameter will tell us how to connect
// to mongodb. This command will fail if the MongoDB server is not available.
//
// Tasks registered in this queue instance will run when `ProcessLoop` is
// called in this *same* instance.
func NewQueue(conf QueueConfig) (monsterqueue.Queue, error) {
	q := &queueMongoDB{
		config: &conf,
		tasks:  make(map[string]monsterqueue.Task),
		done:   make(chan bool),
	}
	var err error
	if conf.Url == "" {
		return nil, errors.New("setting QueueConfig.Url is required")
	}
	dialInfo, err := mgo.ParseURL(conf.Url)
	if err != nil {
		return nil, err
	}
	dialInfo.FailFast = true
	q.session, err = mgo.DialWithInfo(dialInfo)
	if err != nil {
		return nil, err
	}
	q.session.SetSyncTimeout(10 * time.Second)
	q.session.SetSocketTimeout(1 * time.Minute)
	db := q.session.DB(conf.Database)
	if db.Name == "test" {
		q.session.Close()
		return nil, errors.New("database name should be set in QueueConfig.Url or QueueConfig.Database")
	}
	return q, err
}

func (q *queueMongoDB) tasksColl() *mgo.Collection {
	s := q.session.Copy()
	name := "queue_tasks"
	if q.config.CollectionPrefix != "" {
		name = fmt.Sprintf("%s_%s", q.config.CollectionPrefix, name)
	}
	return s.DB(q.config.Database).C(name)
}

func (q *queueMongoDB) RegisterTask(task monsterqueue.Task) error {
	q.tasksMut.Lock()
	defer q.tasksMut.Unlock()
	if _, isRegistered := q.tasks[task.Name()]; isRegistered {
		return errors.New("task already registered")
	}
	q.tasks[task.Name()] = task
	return nil
}

func (q *queueMongoDB) Enqueue(taskName string, params monsterqueue.JobParams) (monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	j := q.initialJob(taskName, params)
	err := coll.Insert(j)
	if err != nil {
		return nil, err
	}
	return &j, nil
}

func (q *queueMongoDB) getDoneJob(jobId bson.ObjectId) (*jobMongoDB, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var resultJob jobMongoDB
	err := coll.Find(bson.M{"_id": jobId, "resultmessage.done": true}).One(&resultJob)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, nil
		}
		return nil, err
	}
	return &resultJob, nil
}

func (q *queueMongoDB) EnqueueWait(taskName string, params monsterqueue.JobParams, timeout time.Duration) (monsterqueue.Job, error) {
	j := q.initialJob(taskName, params)
	j.Waited = true
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	err := coll.Insert(j)
	if err != nil {
		return nil, err
	}
	result := make(chan *jobMongoDB)
	quit := make(chan bool)
	go func() {
		for range time.Tick(200 * time.Millisecond) {
			select {
			case <-quit:
				return
			default:
			}
			job, err := q.getDoneJob(j.Id)
			if err != nil {
				log.Errorf("error trying to get job %s: %s", j.Id, err.Error())
			}
			if job != nil {
				result <- job
				return
			}
		}
	}()
	var resultJob *jobMongoDB
	select {
	case resultJob = <-result:
	case <-time.After(timeout):
		quit <- true
	}
	close(quit)
	close(result)
	if resultJob != nil {
		return resultJob, nil
	}
	err = coll.Update(bson.M{
		"_id":    j.Id,
		"waited": true,
	}, bson.M{"$set": bson.M{"waited": false}})
	if err == mgo.ErrNotFound {
		resultJob, err = q.getDoneJob(j.Id)
	}
	if err != nil {
		return &j, err
	}
	if resultJob != nil {
		return resultJob, nil
	}
	return &j, monsterqueue.ErrQueueWaitTimeout
}

func (q *queueMongoDB) ProcessLoop() {
	for {
		q.wg.Add(1)
		err := q.waitForMessage()
		if err != nil {
			log.Debugf("error getting message from queue: %s", err.Error())
		}
		select {
		case <-time.After(1 * time.Second):
		case <-q.done:
			return
		}
	}
}

func (q *queueMongoDB) Stop() {
	q.done <- true
	q.Wait()
}

func (q *queueMongoDB) Wait() {
	q.wg.Wait()
}

func (q *queueMongoDB) ResetStorage() error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.DropCollection()
}

func (q *queueMongoDB) RetrieveJob(jobId string) (monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var job jobMongoDB
	err := coll.FindId(bson.ObjectIdHex(jobId)).One(&job)
	if err != nil {
		if err == mgo.ErrNotFound {
			return nil, monsterqueue.ErrNoSuchJob
		}
		return nil, err
	}
	return &job, err
}

func (q *queueMongoDB) ListJobs() ([]monsterqueue.Job, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var mongodbJobs []jobMongoDB
	err := coll.Find(nil).All(&mongodbJobs)
	if err != nil {
		return nil, err
	}
	jobs := make([]monsterqueue.Job, len(mongodbJobs))
	for i := range mongodbJobs {
		jobs[i] = &mongodbJobs[i]
	}
	return jobs, nil
}

func (q *queueMongoDB) DeleteJob(jobId string) error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.RemoveId(bson.ObjectIdHex(jobId))
}

func (q *queueMongoDB) initialJob(taskName string, params monsterqueue.JobParams) jobMongoDB {
	buf := make([]byte, monsterqueue.StackTraceLimit)
	buf = buf[:runtime.Stack(buf, false)]
	return jobMongoDB{
		Id:        bson.NewObjectId(),
		Task:      taskName,
		Params:    params,
		Timestamp: time.Now().UTC(),
		Stack:     string(buf),
		queue:     q,
	}
}

func (q *queueMongoDB) waitForMessage() error {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	var job jobMongoDB
	hostname, _ := os.Hostname()
	ownerData := jobOwnership{
		Name:      fmt.Sprintf("%s_%d", hostname, os.Getpid()),
		Owned:     true,
		Timestamp: time.Now().UTC(),
	}
	q.tasksMut.RLock()
	taskNames := make([]string, 0, len(q.tasks))
	for taskName := range q.tasks {
		taskNames = append(taskNames, taskName)
	}
	q.tasksMut.RUnlock()
	_, err := coll.Find(bson.M{
		"task":               bson.M{"$in": taskNames},
		"owner.owned":        false,
		"resultmessage.done": false,
	}).Sort("_id").Apply(mgo.Change{
		Update: bson.M{
			"$set": bson.M{"owner": ownerData},
		},
	}, &job)
	if err != nil {
		q.wg.Done()
		if err == mgo.ErrNotFound {
			return nil
		}
		return err
	}
	job.queue = q
	if err != nil {
		q.moveToResult(&job, nil, err)
		q.wg.Done()
		return err
	}
	q.tasksMut.RLock()
	task, _ := q.tasks[job.Task]
	q.tasksMut.RUnlock()
	if task == nil {
		err := fmt.Errorf("unregistered task name %q", job.Task)
		q.moveToResult(&job, nil, err)
		q.wg.Done()
		return err
	}
	go func() {
		defer q.wg.Done()
		task.Run(&job)
		if !job.ResultMessage.Done {
			q.moveToResult(&job, nil, monsterqueue.ErrNoJobResultSet)
		}
	}()
	return nil
}

func (q *queueMongoDB) moveToResult(job *jobMongoDB, result monsterqueue.JobResult, jobErr error) error {
	var resultMsg jobResultMessage
	resultMsg.Result = result
	resultMsg.Timestamp = time.Now().UTC()
	resultMsg.Done = true
	if jobErr != nil {
		resultMsg.Error = jobErr.Error()
	}
	job.ResultMessage = resultMsg
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	return coll.UpdateId(job.Id, bson.M{"$set": bson.M{"resultmessage": resultMsg, "owner.owned": false}})
}

func (q *queueMongoDB) publishResult(job *jobMongoDB) (bool, error) {
	coll := q.tasksColl()
	defer coll.Database.Session.Close()
	err := coll.Update(bson.M{"_id": job.Id, "waited": true}, bson.M{"$set": bson.M{"waited": false}})
	if err != nil {
		if err == mgo.ErrNotFound {
			return false, nil
		}
		return false, err
	}
	return true, nil
}
<|endoftext|>


<commit_before>package integration

import (
	"testing"

	rbacv1 "k8s.io/api/rbac/v1"
	kapierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kapi "k8s.io/kubernetes/pkg/apis/core"

	authorizationapi "github.com/openshift/origin/pkg/authorization/apis/authorization"
	authorizationclient "github.com/openshift/origin/pkg/authorization/generated/internalclientset"
	testutil "github.com/openshift/origin/test/util"
	testserver "github.com/openshift/origin/test/util/server"
)

func TestRestrictUsers(t *testing.T) {
	masterConfig, err := testserver.DefaultMasterOptions()
	if err != nil {
		t.Fatalf("error creating config: %v", err)
	}
	defer testserver.CleanupMasterEtcd(t, masterConfig)

	masterConfig.KubernetesMasterConfig.APIServerArguments["enable-admission-plugins"] = append(
		masterConfig.KubernetesMasterConfig.APIServerArguments["enable-admission-plugins"],
		"authorization.openshift.io/RestrictSubjectBindings")

	clusterAdminKubeConfig, err := testserver.StartConfiguredMaster(masterConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	clusterAdminKubeClient, err := testutil.GetClusterAdminKubeClient(clusterAdminKubeConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	clusterAdminClientConfig, err := testutil.GetClusterAdminClientConfig(clusterAdminKubeConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	clusterAdminAuthorizationClient := authorizationclient.NewForConfigOrDie(clusterAdminClientConfig).Authorization()

	if _, _, err := testserver.CreateNewProject(clusterAdminClientConfig, "namespace", "carol"); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	role := &authorizationapi.Role{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "role",
		},
	}
	if _, err := clusterAdminAuthorizationClient.Roles("namespace").Create(role); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingAlice := &authorizationapi.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding1",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "alice",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when no restrictions exist should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingAlice); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	allowAlice := &authorizationapi.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-alice",
			Namespace: "namespace",
		},
		Spec: authorizationapi.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationapi.UserRestriction{
				Users: []string{"alice"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowAlice); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingAliceDup := &authorizationapi.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding2",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "alice",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when the subject is already bound should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingAliceDup); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingBob := &authorizationapi.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding3",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "bob",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when the subject is not already bound and is not
	// permitted by any RoleBindingRestrictions should fail.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingBob); !kapierrors.IsForbidden(err) {
		t.Fatalf("expected forbidden, got %v", err)
	}

	// Creating a RBAC rolebinding when the subject is not already bound
	// should also fail.
	rbacRolebindingBob := &rbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding3",
		},
		Subjects: []rbacv1.Subject{
			{
				Kind:      rbacv1.UserKind,
				Namespace: "namespace",
				Name:      "bob",
			},
		},
		RoleRef: rbacv1.RoleRef{Kind: "Role", Name: "role"},
	}
	if _, err := clusterAdminKubeClient.RbacV1().RoleBindings("namespace").Create(rbacRolebindingBob); !kapierrors.IsForbidden(err) {
		t.Fatalf("expected forbidden, got %v", err)
	}

	allowBob := &authorizationapi.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-bob",
			Namespace: "namespace",
		},
		Spec: authorizationapi.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationapi.UserRestriction{
				Users: []string{"bob"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowBob); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Creating a rolebinding when the subject is permitted by some
	// RoleBindingRestrictions should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingBob); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Creating rolebindings that also contains "system non existing" users should
	// not fail.
	allowWithNonExisting := &authorizationapi.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-eve-and-non-existing",
			Namespace: "namespace",
		},
		Spec: authorizationapi.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationapi.UserRestriction{
				Users: []string{"eve", "system:non-existing"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowWithNonExisting); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingEve := &authorizationapi.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding4",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "eve",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingEve); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingNonExisting := &authorizationapi.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding5",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "system:non-existing",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingNonExisting); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}
<commit_msg>integration: use external authorization<commit_after>package integration

import (
	"testing"

	kapi "k8s.io/api/core/v1"
	rbacv1 "k8s.io/api/rbac/v1"
	kapierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	authorizationv1 "github.com/openshift/api/authorization/v1"
	authorizationclient "github.com/openshift/client-go/authorization/clientset/versioned/typed/authorization/v1"
	authorizationapi "github.com/openshift/origin/pkg/authorization/apis/authorization"
	testutil "github.com/openshift/origin/test/util"
	testserver "github.com/openshift/origin/test/util/server"
)

func TestRestrictUsers(t *testing.T) {
	masterConfig, err := testserver.DefaultMasterOptions()
	if err != nil {
		t.Fatalf("error creating config: %v", err)
	}
	defer testserver.CleanupMasterEtcd(t, masterConfig)

	masterConfig.KubernetesMasterConfig.APIServerArguments["enable-admission-plugins"] = append(
		masterConfig.KubernetesMasterConfig.APIServerArguments["enable-admission-plugins"],
		"authorization.openshift.io/RestrictSubjectBindings")
	clusterAdminKubeConfig, err := testserver.StartConfiguredMaster(masterConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	clusterAdminKubeClient, err := testutil.GetClusterAdminKubeClient(clusterAdminKubeConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	clusterAdminClientConfig, err := testutil.GetClusterAdminClientConfig(clusterAdminKubeConfig)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	clusterAdminAuthorizationClient := authorizationclient.NewForConfigOrDie(testutil.NonProtobufConfig(clusterAdminClientConfig))

	if _, _, err := testserver.CreateNewProject(clusterAdminClientConfig, "namespace", "carol"); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	role := &authorizationv1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "role",
		},
	}
	if _, err := clusterAdminAuthorizationClient.Roles("namespace").Create(role); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingAlice := &authorizationv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding1",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "alice",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when no restrictions exist should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingAlice); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	allowAlice := &authorizationv1.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-alice",
			Namespace: "namespace",
		},
		Spec: authorizationv1.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationv1.UserRestriction{
				Users: []string{"alice"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowAlice); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingAliceDup := &authorizationv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding2",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "alice",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when the subject is already bound should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingAliceDup); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingBob := &authorizationv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding3",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "bob",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	// Creating a rolebinding when the subject is not already bound and is not
	// permitted by any RoleBindingRestrictions should fail.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingBob); !kapierrors.IsForbidden(err) {
		t.Fatalf("expected forbidden, got %v", err)
	}

	// Creating a RBAC rolebinding when the subject is not already bound
	// should also fail.
	rbacRolebindingBob := &rbacv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding3",
		},
		Subjects: []rbacv1.Subject{
			{
				Kind:      rbacv1.UserKind,
				Namespace: "namespace",
				Name:      "bob",
			},
		},
		RoleRef: rbacv1.RoleRef{Kind: "Role", Name: "role"},
	}
	if _, err := clusterAdminKubeClient.RbacV1().RoleBindings("namespace").Create(rbacRolebindingBob); !kapierrors.IsForbidden(err) {
		t.Fatalf("expected forbidden, got %v", err)
	}

	allowBob := &authorizationv1.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-bob",
			Namespace: "namespace",
		},
		Spec: authorizationv1.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationv1.UserRestriction{
				Users: []string{"bob"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowBob); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Creating a rolebinding when the subject is permitted by some
	// RoleBindingRestrictions should succeed.
	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingBob); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Creating rolebindings that also contains "system non existing" users should
	// not fail.
	allowWithNonExisting := &authorizationv1.RoleBindingRestriction{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "match-users-eve-and-non-existing",
			Namespace: "namespace",
		},
		Spec: authorizationv1.RoleBindingRestrictionSpec{
			UserRestriction: &authorizationv1.UserRestriction{
				Users: []string{"eve", "system:non-existing"},
			},
		},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindingRestrictions("namespace").Create(allowWithNonExisting); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingEve := &authorizationv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding4",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "eve",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingEve); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rolebindingNonExisting := &authorizationv1.RoleBinding{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: "namespace",
			Name:      "rolebinding5",
		},
		Subjects: []kapi.ObjectReference{
			{
				Kind:      authorizationapi.UserKind,
				Namespace: "namespace",
				Name:      "system:non-existing",
			},
		},
		RoleRef: kapi.ObjectReference{Name: "role", Namespace: "namespace"},
	}

	if _, err := clusterAdminAuthorizationClient.RoleBindings("namespace").Create(rolebindingNonExisting); err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}
<|endoftext|>


<commit_before>// +build integrationtest

package integrationtest

import (
	"fmt"
	"regexp"
	"testing"
	"time"

	crdv1 "github.com/kubernetes-incubator/external-storage/snapshot/pkg/apis/crd/v1"
	"github.com/portworx/sched-ops/k8s"
	"github.com/portworx/torpedo/drivers/scheduler"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
	"k8s.io/apimachinery/pkg/util/wait"
)

var snapRuleFailRegex = regexp.MustCompile("^snapshot failed due to err.+(failed to validate snap rule|failed to run (pre|post)-snap rule).+")

func testSnapshot(t *testing.T) {
	t.Run("simpleSnapshotTest", simpleSnapshotTest)
	t.Run("cloudSnapshotTest", cloudSnapshotTest)
	t.Run("snapshotScaleTest", snapshotScaleTest)
	t.Run("groupSnapshotTest", groupSnapshotTest)
	t.Run("groupSnapshotScaleTest", groupSnapshotScaleTest)
}

func simpleSnapshotTest(t *testing.T) {
	ctx := createSnapshot(t, []string{"mysql-snap-restore"})
	verifySnapshot(t, ctx, "mysql-data", defaultWaitTimeout)
	destroyAndWait(t, ctx)
}

func verifyFailedSnapshot(snapName, snapNamespace string) error {
	failedSnapCheckBackoff := wait.Backoff{
		Duration: 5 * time.Second,
		Factor:   1,
		Steps:    24, // 2 minutes should be enough for the snap to fail
	}

	t := func() (bool, error) {
		snapObj, err := k8s.Instance().GetSnapshot(snapName, snapNamespace)
		if err != nil {
			return false, err
		}

		if snapObj.Status.Conditions == nil {
			return false, nil // conditions not yet populated
		}

		for _, cond := range snapObj.Status.Conditions {
			if cond.Type == crdv1.VolumeSnapshotConditionError {
				if snapRuleFailRegex.MatchString(cond.Message) {
					logrus.Infof("verified that snapshot has failed as expected due to: %s", cond.Message)
					return true, nil
				}
			}
		}

		return false, nil
	}

	return wait.ExponentialBackoff(failedSnapCheckBackoff, t)
}

func cloudSnapshotTest(t *testing.T) {
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: []string{"mysql-cloudsnap-restore"}})
	require.NoError(t, err, "Error scheduling task")
	require.Equal(t, 1, len(ctxs), "Only one task should have started")

	err = schedulerDriver.WaitForRunning(ctxs[0], defaultWaitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for pod to get to running state")

	scheduledNodes, err := schedulerDriver.GetNodesForApp(ctxs[0])
	require.NoError(t, err, "Error getting node for app")
	require.Equal(t, 1, len(scheduledNodes), "App should be scheduled on one node")

	err = schedulerDriver.InspectVolumes(ctxs[0], defaultWaitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for volumes")
	volumeNames := getVolumeNames(t, ctxs[0])
	require.Equal(t, 3, len(volumeNames), "Should only have two volumes and a snapshot")

	dataVolumesNames, dataVolumesInUse := parseDataVolumes(t, "mysql-data", ctxs[0])
	require.Len(t, dataVolumesNames, 2, "should have only 2 data volumes")

	snaps, err := schedulerDriver.GetSnapshots(ctxs[0])
	require.NoError(t, err, "failed to get snapshots")
	require.Len(t, snaps, 1, "should have received exactly one snapshot")

	for _, snap := range snaps {
		s, err := k8s.Instance().GetSnapshot(snap.Name, snap.Namespace)
		require.NoError(t, err, "failed to query snapshot object")
		require.NotNil(t, s, "got nil snapshot object from k8s api")

		require.NotEmpty(t, s.Spec.SnapshotDataName, "snapshot object has empty snapshot data field")

		sData, err := k8s.Instance().GetSnapshotData(s.Spec.SnapshotDataName)
		require.NoError(t, err, "failed to query snapshot data object")

		snapType := sData.Spec.PortworxSnapshot.SnapshotType
		require.Equal(t, snapType, crdv1.PortworxSnapshotTypeCloud)
	}

	fmt.Printf("checking dataVolumesInUse: %v\n", dataVolumesInUse)
	verifyScheduledNode(t, scheduledNodes[0], dataVolumesInUse)
	destroyAndWait(t, ctxs)
}

func groupSnapshotTest(t *testing.T) {
	ctxsToDestroy := make([]*scheduler.Context, 0)
	// Positive tests
	ctxs := createGroupsnaps(t)
	ctxsToDestroy = append(ctxsToDestroy, ctxs...)

	for _, ctx := range ctxs {
		verifyGroupSnapshot(t, ctx, groupSnapshotWaitTimeout)
	}

	// Negative
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: []string{"mysql-snap-group-fail"}})
	require.NoError(t, err, "Error scheduling task")
	require.Len(t, ctxs, 1, "Only one task should have started")

	for _, ctx := range ctxs {
		err = schedulerDriver.WaitForRunning(ctx, defaultWaitTimeout, defaultWaitInterval)
		require.NoError(t, err, "Error waiting for pod to get to running state")

		snaps, err := schedulerDriver.GetSnapshots(ctx)
		require.Error(t, err, "expected to get error when fetching snapahots")
		require.Nil(t, snaps, "expected empty snapshots")
	}
	ctxsToDestroy = append(ctxsToDestroy, ctxs...)

	destroyAndWait(t, ctxsToDestroy)
}

func groupSnapshotScaleTest(t *testing.T) {
	allContexts := make([]*scheduler.Context, 0)
	// Triggers 2 snaps, so use half the count in the loop
	for i := 0; i < snapshotScaleCount/2; i++ {
		ctxs := createGroupsnaps(t)
		allContexts = append(allContexts, ctxs...)
	}

	timeout := groupSnapshotWaitTimeout
	// Increase the timeout if scale is more than or equal 10
	if snapshotScaleCount >= 10 {
		timeout *= time.Duration((snapshotScaleCount / 10) + 1)
	}

	for _, ctx := range allContexts {
		verifyGroupSnapshot(t, ctx, timeout)
	}

	destroyAndWait(t, allContexts)
}

func createGroupsnaps(t *testing.T) []*scheduler.Context {
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: []string{
			"mysql-localsnap-rule",  // tests local group snapshots with a pre exec rule
			"mysql-cloudsnap-group", // tests cloud group snapshots
			"group-cloud-snap-load", // volume is loaded while cloudsnap is being done
		}})
	require.NoError(t, err, "Error scheduling task")
	require.Len(t, ctxs, 3, "Only one task should have started")

	return ctxs
}

func verifyGroupSnapshot(t *testing.T, ctx *scheduler.Context, waitTimeout time.Duration) {
	err := schedulerDriver.WaitForRunning(ctx, waitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for pod to get to running state")

	err = schedulerDriver.InspectVolumes(ctx, waitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error validating storage components")
}

func parseDataVolumes(
	t *testing.T,
	pvcInUseByTest string,
	ctx *scheduler.Context) ([]string, []string) {
	allVolumes, err := schedulerDriver.GetVolumes(ctx)
	require.NoError(t, err, "failed to get volumes")

	dataVolumesNames := make([]string, 0)
	dataVolumesInUse := make([]string, 0)
	for _, v := range allVolumes {
		pvc, err := k8s.Instance().GetPersistentVolumeClaim(v.Name, v.Namespace)
		require.NoError(t, err, "failed to get PVC")

		volName, err := k8s.Instance().GetVolumeForPersistentVolumeClaim(pvc)
		require.NoError(t, err, "failed to get PV name")
		dataVolumesNames = append(dataVolumesNames, volName)

		if pvc.GetName() == pvcInUseByTest {
			dataVolumesInUse = append(dataVolumesInUse, volName)
		}
	}

	require.Len(t, dataVolumesInUse, 1, "should have only 1 data volume in use")

	return dataVolumesNames, dataVolumesInUse
}

func createSnapshot(t *testing.T, appKeys []string) []*scheduler.Context {
	ctx, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: appKeys})
	require.NoError(t, err, "Error scheduling task")
	require.Equal(t, 1, len(ctx), "Only one task should have started")
	return ctx
}

func verifySnapshot(t *testing.T, ctxs []*scheduler.Context, pvcInUseByTest string, waitTimeout time.Duration) {
	err := schedulerDriver.WaitForRunning(ctxs[0], waitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for pod to get to running state")

	scheduledNodes, err := schedulerDriver.GetNodesForApp(ctxs[0])
	require.NoError(t, err, "Error getting node for app")
	require.Equal(t, 1, len(scheduledNodes), "App should be scheduled on one node")

	err = schedulerDriver.InspectVolumes(ctxs[0], waitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for volumes")
	volumeNames := getVolumeNames(t, ctxs[0])
	require.Equal(t, 3, len(volumeNames), "Should only have two volumes and a snapshot")

	dataVolumesNames, dataVolumesInUse := parseDataVolumes(t, pvcInUseByTest, ctxs[0])
	require.Len(t, dataVolumesNames, 2, "should have only 2 data volumes")

	snaps, err := schedulerDriver.GetSnapshots(ctxs[0])
	require.NoError(t, err, "failed to get snapshots")
	require.Len(t, snaps, 1, "should have received exactly one snapshot")

	for _, snap := range snaps {
		s, err := k8s.Instance().GetSnapshot(snap.Name, snap.Namespace)
		require.NoError(t, err, "failed to query snapshot object")
		require.NotNil(t, s, "got nil snapshot object from k8s api")

		require.NotEmpty(t, s.Spec.SnapshotDataName, "snapshot object has empty snapshot data field")

		sData, err := k8s.Instance().GetSnapshotData(s.Spec.SnapshotDataName)
		require.NoError(t, err, "failed to query snapshot data object")

		snapType := sData.Spec.PortworxSnapshot.SnapshotType
		require.Equal(t, snapType, crdv1.PortworxSnapshotTypeLocal)

		snapID := sData.Spec.PortworxSnapshot.SnapshotID
		require.NotEmpty(t, snapID, "got empty snapshot ID in volume snapshot data")

		snapVolInfo, err := storkVolumeDriver.InspectVolume(snapID)
		require.NoError(t, err, "Error getting snapshot volume")
		require.NotNil(t, snapVolInfo.ParentID, "ParentID is nil for snapshot")

		parentVolInfo, err := storkVolumeDriver.InspectVolume(snapVolInfo.ParentID)
		require.NoError(t, err, "Error getting snapshot parent volume")

		parentVolName := parentVolInfo.VolumeName
		var cloneVolName string

		found := false
		for _, volume := range dataVolumesNames {
			if volume == parentVolName {
				found = true
			} else if volume != snapVolInfo.VolumeName {
				cloneVolName = volume
			}
		}
		require.True(t, found, "Parent volume (%v) not found in list of volumes: %v", parentVolName, volumeNames)

		cloneVolInfo, err := storkVolumeDriver.InspectVolume(cloneVolName)
		require.NoError(t, err, "Error getting clone volume")
		require.Equal(t, snapVolInfo.VolumeID, cloneVolInfo.ParentID, "Clone volume does not have snapshot as parent")
	}

	verifyScheduledNode(t, scheduledNodes[0], dataVolumesInUse)
}

func snapshotScaleTest(t *testing.T) {
	ctxs := make([][]*scheduler.Context, snapshotScaleCount)
	for i := 0; i < snapshotScaleCount; i++ {
		ctxs[i] = createSnapshot(t, []string{"mysql-snap-restore"})
	}

	timeout := defaultWaitTimeout
	// Increase the timeout if scale is more than 10
	if snapshotScaleCount > 10 {
		timeout *= time.Duration((snapshotScaleCount / 10) + 1)
	}
	for i := 0; i < snapshotScaleCount; i++ {
		verifySnapshot(t, ctxs[i], "mysql-data", timeout)
	}
	for i := 0; i < snapshotScaleCount; i++ {
		destroyAndWait(t, ctxs[i])
	}
}
<commit_msg>Don't run fio load specs for scale test<commit_after>// +build integrationtest

package integrationtest

import (
	"fmt"
	"regexp"
	"testing"
	"time"

	crdv1 "github.com/kubernetes-incubator/external-storage/snapshot/pkg/apis/crd/v1"
	"github.com/portworx/sched-ops/k8s"
	"github.com/portworx/torpedo/drivers/scheduler"
	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/require"
	"k8s.io/apimachinery/pkg/util/wait"
)

var snapRuleFailRegex = regexp.MustCompile("^snapshot failed due to err.+(failed to validate snap rule|failed to run (pre|post)-snap rule).+")

func testSnapshot(t *testing.T) {
	t.Run("simpleSnapshotTest", simpleSnapshotTest)
	t.Run("cloudSnapshotTest", cloudSnapshotTest)
	t.Run("snapshotScaleTest", snapshotScaleTest)
	t.Run("groupSnapshotTest", groupSnapshotTest)
	t.Run("groupSnapshotScaleTest", groupSnapshotScaleTest)
}

func simpleSnapshotTest(t *testing.T) {
	ctx := createSnapshot(t, []string{"mysql-snap-restore"})
	verifySnapshot(t, ctx, "mysql-data", defaultWaitTimeout)
	destroyAndWait(t, ctx)
}

func verifyFailedSnapshot(snapName, snapNamespace string) error {
	failedSnapCheckBackoff := wait.Backoff{
		Duration: 5 * time.Second,
		Factor:   1,
		Steps:    24, // 2 minutes should be enough for the snap to fail
	}

	t := func() (bool, error) {
		snapObj, err := k8s.Instance().GetSnapshot(snapName, snapNamespace)
		if err != nil {
			return false, err
		}

		if snapObj.Status.Conditions == nil {
			return false, nil // conditions not yet populated
		}

		for _, cond := range snapObj.Status.Conditions {
			if cond.Type == crdv1.VolumeSnapshotConditionError {
				if snapRuleFailRegex.MatchString(cond.Message) {
					logrus.Infof("verified that snapshot has failed as expected due to: %s", cond.Message)
					return true, nil
				}
			}
		}

		return false, nil
	}

	return wait.ExponentialBackoff(failedSnapCheckBackoff, t)
}

func cloudSnapshotTest(t *testing.T) {
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: []string{"mysql-cloudsnap-restore"}})
	require.NoError(t, err, "Error scheduling task")
	require.Equal(t, 1, len(ctxs), "Only one task should have started")

	err = schedulerDriver.WaitForRunning(ctxs[0], defaultWaitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for pod to get to running state")

	scheduledNodes, err := schedulerDriver.GetNodesForApp(ctxs[0])
	require.NoError(t, err, "Error getting node for app")
	require.Equal(t, 1, len(scheduledNodes), "App should be scheduled on one node")

	err = schedulerDriver.InspectVolumes(ctxs[0], defaultWaitTimeout, defaultWaitInterval)
	require.NoError(t, err, "Error waiting for volumes")
	volumeNames := getVolumeNames(t, ctxs[0])
	require.Equal(t, 3, len(volumeNames), "Should only have two volumes and a snapshot")

	dataVolumesNames, dataVolumesInUse := parseDataVolumes(t, "mysql-data", ctxs[0])
	require.Len(t, dataVolumesNames, 2, "should have only 2 data volumes")

	snaps, err := schedulerDriver.GetSnapshots(ctxs[0])
	require.NoError(t, err, "failed to get snapshots")
	require.Len(t, snaps, 1, "should have received exactly one snapshot")

	for _, snap := range snaps {
		s, err := k8s.Instance().GetSnapshot(snap.Name, snap.Namespace)
		require.NoError(t, err, "failed to query snapshot object")
		require.NotNil(t, s, "got nil snapshot object from k8s api")

		require.NotEmpty(t, s.Spec.SnapshotDataName, "snapshot object has empty snapshot data field")

		sData, err := k8s.Instance().GetSnapshotData(s.Spec.SnapshotDataName)
		require.NoError(t, err, "failed to query snapshot data object")

		snapType := sData.Spec.PortworxSnapshot.SnapshotType
		require.Equal(t, snapType, crdv1.PortworxSnapshotTypeCloud)
	}

	fmt.Printf("checking dataVolumesInUse: %v\n", dataVolumesInUse)
	verifyScheduledNode(t, scheduledNodes[0], dataVolumesInUse)
	destroyAndWait(t, ctxs)
}

func groupSnapshotTest(t *testing.T) {
	ctxsToDestroy := make([]*scheduler.Context, 0)
	// Positive tests
	ctxs := createGroupsnaps(t, []string{
		"mysql-localsnap-rule",  // tests local group snapshots with a pre exec rule
		"mysql-cloudsnap-group", // tests cloud group snapshots
		"group-cloud-snap-load", // volume is loaded while cloudsnap is being done
	})

	ctxsToDestroy = append(ctxsToDestroy, ctxs...)

	for _, ctx := range ctxs {
		verifyGroupSnapshot(t, ctx, groupSnapshotWaitTimeout)
	}

	// Negative
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: []string{"mysql-snap-group-fail"}})
	require.NoError(t, err, "Error scheduling task")
	require.Len(t, ctxs, 1, "Only one task should have started")

	for _, ctx := range ctxs {
		err = schedulerDriver.WaitForRunning(ctx, defaultWaitTimeout, defaultWaitInterval)
		require.NoError(t, err, "Error waiting for pod to get to running state")

		snaps, err := schedulerDriver.GetSnapshots(ctx)
		require.Error(t, err, "expected to get error when fetching snapahots")
		require.Nil(t, snaps, "expected empty snapshots")
	}
	ctxsToDestroy = append(ctxsToDestroy, ctxs...)

	destroyAndWait(t, ctxsToDestroy)
}

func groupSnapshotScaleTest(t *testing.T) {
	allContexts := make([]*scheduler.Context, 0)
	// Triggers 2 snaps, so use half the count in the loop
	for i := 0; i < snapshotScaleCount/2; i++ {
		ctxs := createGroupsnaps(t, []string{
			"mysql-localsnap-rule",  // tests local group snapshots with a pre exec rule
			"mysql-cloudsnap-group", // tests cloud group snapshots
		})
		allContexts = append(allContexts, ctxs...)
	}

	timeout := groupSnapshotWaitTimeout
	// Increase the timeout if scale is more than or equal 10
	if snapshotScaleCount >= 10 {
		timeout *= time.Duration((snapshotScaleCount / 10) + 1)
	}

	for _, ctx := range allContexts {
		verifyGroupSnapshot(t, ctx, timeout)
	}

	destroyAndWait(t, allContexts)
}

func createGroupsnaps(t *testing.T, apps []string) []*scheduler.Context {
	ctxs, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: apps})
	require.NoError(t, err, "Error scheduling task")
	require.Len(t, ctxs, len(apps), "Only one task should have started")

	return ctxs
}

func verifyGroupSnapshot(t *testing.T, ctx *scheduler.Context, waitTimeout time.Duration) {
	err := schedulerDriver.WaitForRunning(ctx, waitTimeout, defaultWaitInterval)
	require.NoError(t, err, fmt.Sprintf("Error waiting for app to get to running state in context: %s-%s", ctx.App.Key, ctx.UID))

	err = schedulerDriver.InspectVolumes(ctx, waitTimeout, defaultWaitInterval)
	require.NoError(t, err, fmt.Sprintf("Error validating storage components in context: %s-%s", ctx.App.Key, ctx.UID))
}

func parseDataVolumes(
	t *testing.T,
	pvcInUseByTest string,
	ctx *scheduler.Context) ([]string, []string) {
	allVolumes, err := schedulerDriver.GetVolumes(ctx)
	require.NoError(t, err, "failed to get volumes")

	dataVolumesNames := make([]string, 0)
	dataVolumesInUse := make([]string, 0)
	for _, v := range allVolumes {
		pvc, err := k8s.Instance().GetPersistentVolumeClaim(v.Name, v.Namespace)
		require.NoError(t, err, "failed to get PVC")

		volName, err := k8s.Instance().GetVolumeForPersistentVolumeClaim(pvc)
		require.NoError(t, err, "failed to get PV name")
		dataVolumesNames = append(dataVolumesNames, volName)

		if pvc.GetName() == pvcInUseByTest {
			dataVolumesInUse = append(dataVolumesInUse, volName)
		}
	}

	require.Len(t, dataVolumesInUse, 1, "should have only 1 data volume in use")

	return dataVolumesNames, dataVolumesInUse
}

func createSnapshot(t *testing.T, appKeys []string) []*scheduler.Context {
	ctx, err := schedulerDriver.Schedule(generateInstanceID(t, ""),
		scheduler.ScheduleOptions{AppKeys: appKeys})
	require.NoError(t, err, "Error scheduling task")
	require.Equal(t, 1, len(ctx), "Only one task should have started")
	return ctx
}

func verifySnapshot(t *testing.T, ctxs []*scheduler.Context, pvcInUseByTest string, waitTimeout time.Duration) {
	err := schedulerDriver.WaitForRunning(ctxs[0], waitTimeout, defaultWaitInterval)
	require.NoError(t, err, fmt.Sprintf("Error waiting for app to get to running state in context: %s-%s", ctxs[0].App.Key, ctxs[0].UID))

	scheduledNodes, err := schedulerDriver.GetNodesForApp(ctxs[0])
	require.NoError(t, err, "Error getting node for app")
	require.Equal(t, 1, len(scheduledNodes), "App should be scheduled on one node")

	err = schedulerDriver.InspectVolumes(ctxs[0], waitTimeout, defaultWaitInterval)
	require.NoError(t, err, fmt.Sprintf("Error waiting for volumes in context: %s-%s", ctxs[0].App.Key, ctxs[0].UID))
	volumeNames := getVolumeNames(t, ctxs[0])
	require.Equal(t, 3, len(volumeNames), "Should only have two volumes and a snapshot")

	dataVolumesNames, dataVolumesInUse := parseDataVolumes(t, pvcInUseByTest, ctxs[0])
	require.Len(t, dataVolumesNames, 2, "should have only 2 data volumes")

	snaps, err := schedulerDriver.GetSnapshots(ctxs[0])
	require.NoError(t, err, "failed to get snapshots")
	require.Len(t, snaps, 1, "should have received exactly one snapshot")

	for _, snap := range snaps {
		s, err := k8s.Instance().GetSnapshot(snap.Name, snap.Namespace)
		require.NoError(t, err, "failed to query snapshot object")
		require.NotNil(t, s, "got nil snapshot object from k8s api")

		require.NotEmpty(t, s.Spec.SnapshotDataName, "snapshot object has empty snapshot data field")

		sData, err := k8s.Instance().GetSnapshotData(s.Spec.SnapshotDataName)
		require.NoError(t, err, "failed to query snapshot data object")

		snapType := sData.Spec.PortworxSnapshot.SnapshotType
		require.Equal(t, snapType, crdv1.PortworxSnapshotTypeLocal)

		snapID := sData.Spec.PortworxSnapshot.SnapshotID
		require.NotEmpty(t, snapID, "got empty snapshot ID in volume snapshot data")

		snapVolInfo, err := storkVolumeDriver.InspectVolume(snapID)
		require.NoError(t, err, "Error getting snapshot volume")
		require.NotNil(t, snapVolInfo.ParentID, "ParentID is nil for snapshot")

		parentVolInfo, err := storkVolumeDriver.InspectVolume(snapVolInfo.ParentID)
		require.NoError(t, err, "Error getting snapshot parent volume")

		parentVolName := parentVolInfo.VolumeName
		var cloneVolName string

		found := false
		for _, volume := range dataVolumesNames {
			if volume == parentVolName {
				found = true
			} else if volume != snapVolInfo.VolumeName {
				cloneVolName = volume
			}
		}
		require.True(t, found, "Parent volume (%v) not found in list of volumes: %v", parentVolName, volumeNames)

		cloneVolInfo, err := storkVolumeDriver.InspectVolume(cloneVolName)
		require.NoError(t, err, "Error getting clone volume")
		require.Equal(t, snapVolInfo.VolumeID, cloneVolInfo.ParentID, "Clone volume does not have snapshot as parent")
	}

	verifyScheduledNode(t, scheduledNodes[0], dataVolumesInUse)
}

func snapshotScaleTest(t *testing.T) {
	ctxs := make([][]*scheduler.Context, snapshotScaleCount)
	for i := 0; i < snapshotScaleCount; i++ {
		ctxs[i] = createSnapshot(t, []string{"mysql-snap-restore"})
	}

	timeout := defaultWaitTimeout
	// Increase the timeout if scale is more than 10
	if snapshotScaleCount > 10 {
		timeout *= time.Duration((snapshotScaleCount / 10) + 1)
	}
	for i := 0; i < snapshotScaleCount; i++ {
		verifySnapshot(t, ctxs[i], "mysql-data", timeout)
	}
	for i := 0; i < snapshotScaleCount; i++ {
		destroyAndWait(t, ctxs[i])
	}
}
<|endoftext|>