예제 #1
0
func TestKinesumerBeginEnd(t *testing.T) {
	k, kin, sssm, prov := makeTestKinesumer(t)
	k.Stream = "c"

	kin.On("DescribeStreamPages", mock.Anything, mock.Anything).Return(awserr.New("bad", "bad", nil)).Once()
	_, err := k.Begin()
	assert.Error(t, err)

	prov.On("TTL").Return(time.Millisecond * 10)
	prov.On("TryAcquire", mock.Anything).Return(nil)
	prov.On("Heartbeat", mock.Anything).Return(nil)
	prov.On("Release", mock.Anything).Return(nil)
	kin.On("DescribeStreamPages", mock.Anything, mock.Anything).Return(awserr.Error(nil))
	sssm.On("Begin", mock.Anything).Return(nil)
	sssm.On("GetStartSequence", mock.Anything).Return("0").Once()
	sssm.On("GetStartSequence", mock.Anything).Return("")
	sssm.On("TryAcquire", mock.Anything).Return(nil)
	kin.On("GetShardIterator", mock.Anything).Return(&kinesis.GetShardIteratorOutput{
		ShardIterator: aws.String("0"),
	}, awserr.Error(nil))
	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAAA"),
		Records:            []*kinesis.Record{},
	}, awserr.Error(nil))
	sssm.On("End").Return()
	_, err = k.Begin()
	assert.Nil(t, err)
	assert.Equal(t, 2, k.nRunning)
	k.End()
}
예제 #2
0
func TestShardWorkerGetShardIterator(t *testing.T) {
	s, kin, _, _, _, _, _ := makeTestShardWorker()

	kin.On("GetShardIterator", mock.Anything).Return(&kinesis.GetShardIteratorOutput{
		ShardIterator: aws.String("AAAAA"),
	}, awserr.Error(nil))
	res, err := s.GetShardIterator("TYPE", "123")
	assert.Nil(t, err)
	assert.Equal(t, "AAAAA", res)
}
예제 #3
0
func TestShardWorkerGetRecordsAndProcess(t *testing.T) {
	s, kin, sssm, prov, stp, _, c := makeTestShardWorker()

	prov.On("Heartbeat", mock.Anything).Return(nil)

	record1 := kinesis.Record{
		Data:           []byte("help I'm trapped"),
		PartitionKey:   aws.String("aaaa"),
		SequenceNumber: aws.String("123"),
	}
	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAA"),
		Records:            []*kinesis.Record{&record1},
	}, awserr.Error(nil)).Once()
	doneC := make(chan k.Record)
	sssm.On("DoneC").Return(doneC)
	brk, nextIt, nextSeq := s.GetRecordsAndProcess("AAAA", "123")
	rec := <-c
	assert.Equal(t, record1.Data, rec.Data())
	assert.False(t, brk)
	assert.Equal(t, "AAAA", nextIt)
	assert.Equal(t, "123", nextSeq)

	err := awserr.New("bad", "bad", nil)
	stp <- Unit{}
	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAA"),
		Records:            []*kinesis.Record{},
	}, err)
	kin.On("GetShardIterator", mock.Anything).Return(&kinesis.GetShardIteratorOutput{
		ShardIterator: aws.String("AAAA"),
	}, awserr.Error(nil))
	brk, nextIt, nextSeq = s.GetRecordsAndProcess("AAAA", "123")
	kin.AssertNumberOfCalls(t, "GetShardIterator", 1)
	assert.True(t, brk)
}
예제 #4
0
func (input *Sqs3Input) Run(runner pipeline.InputRunner,
    helper pipeline.PluginHelper) error {

    // initialize
    input.runner = runner
    input.sqs = sqs.New(session.New())
    input.s3 = s3.New(session.New())
    queue_url, err := get_queue(input.sqs, input.SqsQueue)
    if err != nil { return err }
    input.queue_url = queue_url
    //input.hostname = helper.PipelineConfig().Hostname()
    tickChan := runner.Ticker()
    sRunner := runner.NewSplitterRunner("")
    if !sRunner.UseMsgBytes() {
        sRunner.SetPackDecorator(input.packDecorator)
    }
    defer sRunner.Done()

    for {
        select {
        case <-input.stop:
            return nil
        case <-tickChan:
        }

        receipt_handle, bucket, key, err := receive_from_queue(input.sqs, input.queue_url)
        if err != nil {
            runner.LogError(fmt.Errorf("Error reading queue: %s", err.Error()))
            continue
        }

        o, _, err := get_object(input.s3, bucket, key)
        if err != nil {
            runner.LogError(fmt.Errorf("Error opening file: %s", err.Error()))
            if aws_err := awserr.Error(err); aws_err != nil {
                f aws_err.Code == "NoSuchBucket" or aws_err.Code == "NoSuchKey" {
                    delete_message(input.sqs, input.queue_url, receipt_handle)
                }
            }
            continue
        }
        for err == nil {
            err = sRunner.SplitStream(o, nil)
            if err != io.EOF && err != nil {
                runner.LogError(fmt.Errorf("Error reading file: %s", err.Error()))
            }
        }
        o.Close()
    }
}
예제 #5
0
func TestShardWorkerGetRecords(t *testing.T) {
	s, kin, _, _, _, _, _ := makeTestShardWorker()

	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAA"),
		Records:            []*kinesis.Record{},
	}, awserr.Error(nil))

	records, nextIt, mills, err := s.GetRecords("AAAA")
	assert.Nil(t, err)
	assert.Equal(t, 0, len(records))
	assert.Equal(t, "AAAA", nextIt)
	assert.Equal(t, int64(0), mills)
}
예제 #6
0
func TestShardWorkerRun(t *testing.T) {
	s, kin, sssm, prov, stp, stpd, c := makeTestShardWorker()

	prov.On("Heartbeat", mock.Anything).Return(nil)
	prov.On("Release", mock.Anything).Return(nil)
	sssm.On("GetStartSequence", mock.Anything).Return("AAAA")

	record1 := kinesis.Record{
		Data:           []byte("help I'm trapped"),
		PartitionKey:   aws.String("aaaa"),
		SequenceNumber: aws.String("123"),
	}
	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAA"),
		Records:            []*kinesis.Record{&record1},
	}, awserr.Error(nil)).Once()
	kin.On("GetRecords", mock.Anything).Return(&kinesis.GetRecordsOutput{
		MillisBehindLatest: aws.Int64(0),
		NextShardIterator:  aws.String("AAAA"),
		Records:            []*kinesis.Record{},
	}, awserr.Error(nil))
	doneC := make(chan k.Record)
	sssm.On("DoneC").Return(doneC)
	kin.On("GetShardIterator", mock.Anything).Return(&kinesis.GetShardIteratorOutput{
		ShardIterator: aws.String("AAAA"),
	}, awserr.Error(nil))
	go func() {
		time.Sleep(10 * time.Millisecond)
		stp <- Unit{}
	}()
	s.RunWorker()
	<-stpd
	rec := <-c
	assert.Equal(t, record1.Data, rec.Data())
}