/* * Copyright (c) 2018 VMware, Inc. * * Permission is hereby granted, free of charge, to any person obtaining a copy of this software and * associated documentation files (the "Software"), to deal in the Software without restriction, including * without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell * copies of the Software, and to permit persons to whom the Software is furnished to do * so, subject to the following conditions: * * The above copyright notice and this permission notice shall be included in all copies or substantial * portions of the Software. * * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT * NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ package worker import ( "net/http" "os" "testing" "time" "github.com/aws/aws-sdk-go/aws" "github.com/prometheus/common/expfmt" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" cfg "vmware.com/cascade-kinesis-client/clientlibrary/config" kc "vmware.com/cascade-kinesis-client/clientlibrary/interfaces" "vmware.com/cascade-kinesis-client/clientlibrary/metrics" "vmware.com/cascade-kinesis-client/clientlibrary/utils" ) const ( streamName = "kcl-test" regionName = "us-west-2" workerID = "test-worker" ) const specstr = `{"name":"kube-qQyhk","networking":{"containerNetworkCidr":"10.2.0.0/16"},"orgName":"BVT-Org-cLQch","projectName":"project-tDSJd","serviceLevel":"DEVELOPER","size":{"count":1},"version":"1.8.1-4"}` const metricsSystem = "cloudwatch" func TestWorker(t *testing.T) { kclConfig := cfg.NewKinesisClientLibConfig("appName", streamName, regionName, workerID). WithInitialPositionInStream(cfg.LATEST). WithMaxRecords(10). WithMaxLeasesForWorker(1). WithShardSyncIntervalMillis(5000). WithFailoverTimeMillis(300000) log.SetOutput(os.Stdout) log.SetLevel(log.DebugLevel) assert.Equal(t, regionName, kclConfig.RegionName) assert.Equal(t, streamName, kclConfig.StreamName) // configure cloudwatch as metrics system metricsConfig := getMetricsConfig(metricsSystem) worker := NewWorker(recordProcessorFactory(t), kclConfig, metricsConfig) assert.Equal(t, regionName, worker.regionName) assert.Equal(t, streamName, worker.streamName) err := worker.Start() assert.Nil(t, err) // Put some data into stream. for i := 0; i < 100; i++ { // Use random string as partition key to ensure even distribution across shards err := worker.Publish(streamName, utils.RandStringBytesMaskImpr(10), []byte(specstr)) if err != nil { t.Errorf("Errorin Publish. %+v", err) } } // wait a few seconds before shutdown processing time.Sleep(10 * time.Second) if metricsConfig != nil && metricsConfig.MonitoringService == "prometheus" { res, err := http.Get("http://localhost:8080/metrics") if err != nil { t.Fatalf("Error scraping Prometheus endpoint %s", err) } var parser expfmt.TextParser parsed, err := parser.TextToMetricFamilies(res.Body) res.Body.Close() if err != nil { t.Errorf("Error reading monitoring response %s", err) } t.Logf("Prometheus: %+v", parsed) } worker.Shutdown() } // configure different metrics system func getMetricsConfig(service string) *metrics.MonitoringConfiguration { if service == "cloudwatch" { return &metrics.MonitoringConfiguration{ MonitoringService: "cloudwatch", Region: regionName, CloudWatch: metrics.CloudWatchMonitoringService{ // Those value should come from kclConfig MetricsBufferTimeMillis: 10000, MetricsMaxQueueSize: 20, }, } } if service == "prometheus" { return &metrics.MonitoringConfiguration{ MonitoringService: "prometheus", Region: regionName, Prometheus: metrics.PrometheusMonitoringService{ ListenAddress: ":8080", }, } } return nil } // Record processor factory is used to create RecordProcessor func recordProcessorFactory(t *testing.T) kc.IRecordProcessorFactory { return &dumpRecordProcessorFactory{t: t} } // simple record processor and dump everything type dumpRecordProcessorFactory struct { t *testing.T } func (d *dumpRecordProcessorFactory) CreateProcessor() kc.IRecordProcessor { return &dumpRecordProcessor{ t: d.t, } } // Create a dump record processor for printing out all data from record. type dumpRecordProcessor struct { t *testing.T } func (dd *dumpRecordProcessor) Initialize(input *kc.InitializationInput) { dd.t.Logf("Processing SharId: %v at checkpoint: %v", input.ShardId, aws.StringValue(input.ExtendedSequenceNumber.SequenceNumber)) } func (dd *dumpRecordProcessor) ProcessRecords(input *kc.ProcessRecordsInput) { dd.t.Log("Processing Records...") // don't process empty record if len(input.Records) == 0 { return } for _, v := range input.Records { dd.t.Logf("Record = %s", v.Data) assert.Equal(dd.t, specstr, string(v.Data)) } // checkpoint it after processing this batch lastRecordSequenceNubmer := input.Records[len(input.Records)-1].SequenceNumber dd.t.Logf("Checkpoint progress at: %v, MillisBehindLatest = %v", lastRecordSequenceNubmer, input.MillisBehindLatest) input.Checkpointer.Checkpoint(lastRecordSequenceNubmer) } func (dd *dumpRecordProcessor) Shutdown(input *kc.ShutdownInput) { dd.t.Logf("Shutdown Reason: %v", aws.StringValue(kc.ShutdownReasonMessage(input.ShutdownReason))) // When the value of {@link ShutdownInput#getShutdownReason()} is // {@link com.amazonaws.services.kinesis.clientlibrary.lib.worker.ShutdownReason#TERMINATE} it is required that you // checkpoint. Failure to do so will result in an IllegalArgumentException, and the KCL no longer making progress. if input.ShutdownReason == kc.TERMINATE { input.Checkpointer.Checkpoint(nil) } }