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()
}
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)
}
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)
}
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()
}
}
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)
}
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())
}
