From f84b03f2528702de85273cc891b12cae3ade3800 Mon Sep 17 00:00:00 2001 From: sundowndev Date: Tue, 1 Feb 2022 15:26:29 +0400 Subject: [PATCH] feat: add field to telemetry --- pkg/driftctl.go | 5 +- pkg/driftctl_test.go | 13 ++++- pkg/iac/supplier/IacChainSupplier.go | 12 ++++- pkg/iac/supplier/IacChainSupplier_test.go | 20 ++++---- pkg/iac/supplier/supplier.go | 4 +- .../terraform/state/terraform_state_reader.go | 17 ++++++- pkg/resource/mock_IaCSupplier.go | 47 +++++++++++++++++++ pkg/resource/supplier.go | 8 +++- pkg/telemetry/telemetry.go | 5 ++ pkg/telemetry/telemetry_test.go | 2 + 10 files changed, 114 insertions(+), 19 deletions(-) create mode 100644 pkg/resource/mock_IaCSupplier.go diff --git a/pkg/driftctl.go b/pkg/driftctl.go index 57cfc7b07..8032249f1 100644 --- a/pkg/driftctl.go +++ b/pkg/driftctl.go @@ -38,7 +38,7 @@ type ScanOptions struct { type DriftCTL struct { remoteSupplier resource.Supplier - iacSupplier resource.Supplier + iacSupplier resource.IaCSupplier alerter alerter.AlerterInterface analyzer *analyser.Analyzer resourceFactory resource.ResourceFactory @@ -50,7 +50,7 @@ type DriftCTL struct { } func NewDriftCTL(remoteSupplier resource.Supplier, - iacSupplier resource.Supplier, + iacSupplier resource.IaCSupplier, alerter *alerter.Alerter, analyzer *analyser.Analyzer, resFactory resource.ResourceFactory, @@ -189,6 +189,7 @@ func (d DriftCTL) scan() (remoteResources []*resource.Resource, resourcesFromSta if err != nil { return nil, nil, err } + d.store.Bucket(memstore.TelemetryBucket).Set("iac_source_count", d.iacSupplier.SourceCount()) logrus.Info("Start scanning cloud provider") d.scanProgress.Start() diff --git a/pkg/driftctl_test.go b/pkg/driftctl_test.go index a8b798bb3..f724e276d 100644 --- a/pkg/driftctl_test.go +++ b/pkg/driftctl_test.go @@ -62,8 +62,9 @@ func runTest(t *testing.T, cases TestCases) { res.Sch = schema } - stateSupplier := &resource.MockSupplier{} + stateSupplier := &resource.MockIaCSupplier{} stateSupplier.On("Resources").Return(c.stateResources, nil) + stateSupplier.On("SourceCount").Return(uint(2)) if c.remoteResources == nil { c.remoteResources = []*resource.Resource{} @@ -138,6 +139,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -155,6 +157,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, options: func(t *testing.T) *pkg.ScanOptions { return &pkg.ScanOptions{Deep: true} @@ -173,6 +176,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -226,6 +230,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -264,6 +269,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -304,6 +310,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -344,6 +351,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, }, { @@ -434,6 +442,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 4, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 2, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, options: func(t *testing.T) *pkg.ScanOptions { return &pkg.ScanOptions{ @@ -530,6 +539,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 6, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 2, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, options: func(t *testing.T) *pkg.ScanOptions { return &pkg.ScanOptions{ @@ -628,6 +638,7 @@ func TestDriftctlRun_BasicBehavior(t *testing.T) { assert.Equal(t, 1, store.Bucket(memstore.TelemetryBucket).Get("total_resources")) assert.Equal(t, 0, store.Bucket(memstore.TelemetryBucket).Get("total_managed")) assert.Equal(t, uint(0), store.Bucket(memstore.TelemetryBucket).Get("duration")) + assert.Equal(t, uint(2), store.Bucket(memstore.TelemetryBucket).Get("iac_source_count")) }, options: func(t *testing.T) *pkg.ScanOptions { filterStr := "Id=='role-test-1'" diff --git a/pkg/iac/supplier/IacChainSupplier.go b/pkg/iac/supplier/IacChainSupplier.go index 6dc94e90d..d63d5461f 100644 --- a/pkg/iac/supplier/IacChainSupplier.go +++ b/pkg/iac/supplier/IacChainSupplier.go @@ -10,7 +10,7 @@ import ( ) type IacChainSupplier struct { - suppliers []resource.Supplier + suppliers []resource.IaCSupplier runner *parallel.ParallelRunner } @@ -20,7 +20,15 @@ func NewIacChainSupplier() *IacChainSupplier { } } -func (r *IacChainSupplier) AddSupplier(supplier resource.Supplier) { +func (r *IacChainSupplier) SourceCount() uint { + count := uint(0) + for _, supplier := range r.suppliers { + count += supplier.SourceCount() + } + return count +} + +func (r *IacChainSupplier) AddSupplier(supplier resource.IaCSupplier) { r.suppliers = append(r.suppliers, supplier) } diff --git a/pkg/iac/supplier/IacChainSupplier_test.go b/pkg/iac/supplier/IacChainSupplier_test.go index 64def7ce4..9c7d7885b 100644 --- a/pkg/iac/supplier/IacChainSupplier_test.go +++ b/pkg/iac/supplier/IacChainSupplier_test.go @@ -11,22 +11,22 @@ import ( func TestIacChainSupplier_Resources(t *testing.T) { tests := []struct { name string - initSuppliers func(suppliers *[]resource.Supplier) + initSuppliers func(suppliers *[]resource.IaCSupplier) want []*resource.Resource wantErr bool }{ { name: "All failed", - initSuppliers: func(suppliers *[]resource.Supplier) { - sup := &resource.MockSupplier{} + initSuppliers: func(suppliers *[]resource.IaCSupplier) { + sup := &resource.MockIaCSupplier{} sup.On("Resources").Return(nil, errors.New("1")) *suppliers = append(*suppliers, sup) - sup = &resource.MockSupplier{} + sup = &resource.MockIaCSupplier{} sup.On("Resources").Return(nil, errors.New("2")) *suppliers = append(*suppliers, sup) - sup = &resource.MockSupplier{} + sup = &resource.MockIaCSupplier{} sup.On("Resources").Return(nil, errors.New("3")) *suppliers = append(*suppliers, sup) }, @@ -35,16 +35,16 @@ func TestIacChainSupplier_Resources(t *testing.T) { }, { name: "Partial failed", - initSuppliers: func(suppliers *[]resource.Supplier) { - sup := &resource.MockSupplier{} + initSuppliers: func(suppliers *[]resource.IaCSupplier) { + sup := &resource.MockIaCSupplier{} sup.On("Resources").Return(nil, errors.New("1")) *suppliers = append(*suppliers, sup) - sup = &resource.MockSupplier{} + sup = &resource.MockIaCSupplier{} sup.On("Resources").Return(nil, errors.New("2")) *suppliers = append(*suppliers, sup) - sup = &resource.MockSupplier{} + sup = &resource.MockIaCSupplier{} sup.On("Resources").Return([]*resource.Resource{ &resource.Resource{ Id: "ID", @@ -67,7 +67,7 @@ func TestIacChainSupplier_Resources(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { r := NewIacChainSupplier() - suppliers := make([]resource.Supplier, 0) + suppliers := make([]resource.IaCSupplier, 0) tt.initSuppliers(&suppliers) for _, supplier := range suppliers { diff --git a/pkg/iac/supplier/supplier.go b/pkg/iac/supplier/supplier.go index f21a376c8..184f794d4 100644 --- a/pkg/iac/supplier/supplier.go +++ b/pkg/iac/supplier/supplier.go @@ -37,7 +37,7 @@ func GetIACSupplier(configs []config.SupplierConfig, progress output.Progress, alerter *alerter.Alerter, factory resource.ResourceFactory, - filter filter.Filter) (resource.Supplier, error) { + filter filter.Filter) (resource.IaCSupplier, error) { chainSupplier := NewIacChainSupplier() for _, config := range configs { @@ -47,7 +47,7 @@ func GetIACSupplier(configs []config.SupplierConfig, deserializer := resource.NewDeserializer(factory) - var supplier resource.Supplier + var supplier resource.IaCSupplier var err error switch config.Key { case state.TerraformStateReaderSupplier: diff --git a/pkg/iac/terraform/state/terraform_state_reader.go b/pkg/iac/terraform/state/terraform_state_reader.go index 11836e208..2cf344af3 100644 --- a/pkg/iac/terraform/state/terraform_state_reader.go +++ b/pkg/iac/terraform/state/terraform_state_reader.go @@ -41,6 +41,7 @@ type TerraformStateReader struct { progress output.Progress filter filter.Filter alerter *alerter.Alerter + sourceCount uint } func (r *TerraformStateReader) initReader() error { @@ -49,7 +50,16 @@ func (r *TerraformStateReader) initReader() error { } func NewReader(config config.SupplierConfig, library *terraform.ProviderLibrary, backendOpts *backend.Options, progress output.Progress, alerter *alerter.Alerter, deserializer *resource.Deserializer, filter filter.Filter) (*TerraformStateReader, error) { - reader := TerraformStateReader{library: library, config: config, deserializer: deserializer, backendOptions: backendOpts, progress: progress, alerter: alerter, filter: filter} + reader := TerraformStateReader{ + library: library, + config: config, + deserializer: deserializer, + backendOptions: backendOpts, + progress: progress, + alerter: alerter, + filter: filter, + sourceCount: 0, + } err := reader.initReader() if err != nil { return nil, err @@ -214,8 +224,13 @@ func (r *TerraformStateReader) Resources() ([]*resource.Resource, error) { return r.retrieveMultiplesStates() } +func (r *TerraformStateReader) SourceCount() uint { + return r.sourceCount +} + func (r *TerraformStateReader) retrieveForState(path string) ([]*resource.Resource, error) { r.config.Path = path + r.sourceCount += 1 logrus.WithFields(logrus.Fields{ "path": r.config.Path, "backend": r.config.Backend, diff --git a/pkg/resource/mock_IaCSupplier.go b/pkg/resource/mock_IaCSupplier.go new file mode 100644 index 000000000..30907b86d --- /dev/null +++ b/pkg/resource/mock_IaCSupplier.go @@ -0,0 +1,47 @@ +// Code generated by mockery v0.0.0-dev. DO NOT EDIT. + +package resource + +import mock "github.com/stretchr/testify/mock" + +// MockIaCSupplier is an autogenerated mock type for the IaCSupplier type +type MockIaCSupplier struct { + mock.Mock +} + +// Resources provides a mock function with given fields: +func (_m *MockIaCSupplier) Resources() ([]*Resource, error) { + ret := _m.Called() + + var r0 []*Resource + if rf, ok := ret.Get(0).(func() []*Resource); ok { + r0 = rf() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).([]*Resource) + } + } + + var r1 error + if rf, ok := ret.Get(1).(func() error); ok { + r1 = rf() + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +// SourceCount provides a mock function with given fields: +func (_m *MockIaCSupplier) SourceCount() uint { + ret := _m.Called() + + var r0 uint + if rf, ok := ret.Get(0).(func() uint); ok { + r0 = rf() + } else { + r0 = ret.Get(0).(uint) + } + + return r0 +} diff --git a/pkg/resource/supplier.go b/pkg/resource/supplier.go index 626a239d6..070fc8da0 100644 --- a/pkg/resource/supplier.go +++ b/pkg/resource/supplier.go @@ -1,10 +1,16 @@ package resource -// Resource Supplier supply the list of resource.Resource, its the front to retrieve remote resources +// Supplier supply the list of resource.Resource, it's the main interface to retrieve remote resources type Supplier interface { Resources() ([]*Resource, error) } +// IaCSupplier supply the list of resource.Resource, it's the main interface to retrieve state resources +type IaCSupplier interface { + Supplier + SourceCount() uint +} + type StoppableSupplier interface { Supplier Stop() diff --git a/pkg/telemetry/telemetry.go b/pkg/telemetry/telemetry.go index 33e9c47e2..1102e86b3 100644 --- a/pkg/telemetry/telemetry.go +++ b/pkg/telemetry/telemetry.go @@ -20,6 +20,7 @@ type telemetry struct { TotalManaged int `json:"total_managed"` Duration uint `json:"duration"` ProviderName string `json:"provider_name"` + IaCSourceCount uint `json:"iac_source_count"` } type Telemetry struct { @@ -59,6 +60,10 @@ func (te Telemetry) SendTelemetry(store memstore.Bucket) { t.ProviderName = val } + if val, ok := store.Get("iac_source_count").(uint); ok { + t.IaCSourceCount = val + } + body, err := json.Marshal(t) if err != nil { logrus.Debug(err) diff --git a/pkg/telemetry/telemetry_test.go b/pkg/telemetry/telemetry_test.go index 5ccc74891..51c400eb7 100644 --- a/pkg/telemetry/telemetry_test.go +++ b/pkg/telemetry/telemetry_test.go @@ -44,12 +44,14 @@ func TestSendTelemetry(t *testing.T) { TotalManaged: 1, Duration: 123, ProviderName: "aws", + IaCSourceCount: 2, }, setStoreValues: func(s memstore.Bucket, a *analyser.Analysis) { s.Set("total_resources", a.Summary().TotalResources) s.Set("total_managed", a.Summary().TotalManaged) s.Set("duration", uint(a.Duration.Seconds()+0.5)) s.Set("provider_name", "aws") + s.Set("iac_source_count", uint(2)) }, }, {