Fix mocks

elastic · Mar 11, 2024 · 14d57a4 · 14d57a4
1 parent 9e06910
commit 14d57a4
Show file tree

Hide file tree

Showing 6 changed files with 411 additions and 77 deletions.
diff --git a/internal/inventory/aws/fetcher_ec2_test.go b/internal/inventory/aws/fetcher_ec2_test.go
@@ -25,21 +25,13 @@ import (
 	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/elastic/elastic-agent-libs/logp"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/elastic/cloudbeat/internal/inventory"
 	ec2beat "github.com/elastic/cloudbeat/internal/resources/providers/awslib/ec2"
 	"github.com/elastic/cloudbeat/internal/resources/utils/pointers"
 )
 
-type mockInstancesProvider struct {
-	instances []*ec2beat.Ec2Instance
-	err       error
-}
-
-func (m *mockInstancesProvider) DescribeInstances(_ context.Context) ([]*ec2beat.Ec2Instance, error) {
-	return m.instances, m.err
-}
-
 func TestFetch(t *testing.T) {
 	instance1 := &ec2beat.Ec2Instance{
 		Instance: types.Instance{
@@ -130,11 +122,12 @@ func TestFetch(t *testing.T) {
 	}
 
 	logger := logp.NewLogger("test_fetcher_ec2")
+	provider := newMockInstancesProvider(t)
+	provider.EXPECT().DescribeInstances(mock.Anything).Return(in, nil)
+
 	fetcher := Ec2Fetcher{
-		logger: logger,
-		provider: &mockInstancesProvider{
-			instances: in,
-		},
+		logger:   logger,
+		provider: provider,
 	}
 
 	ch := make(chan inventory.AssetEvent)

diff --git a/internal/inventory/aws/mock_instances_provider.go b/internal/inventory/aws/mock_instances_provider.go
diff --git a/internal/inventory/inventory_test.go b/internal/inventory/inventory_test.go
@@ -28,28 +28,11 @@ import (
 	"github.com/elastic/elastic-agent-libs/logp"
 	"github.com/elastic/elastic-agent-libs/mapstr"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/elastic/cloudbeat/internal/resources/utils/pointers"
 )
 
-type mockAssetPublisher struct {
-	mockF func(events []beat.Event)
-}
-
-func (m mockAssetPublisher) PublishAll(events []beat.Event) {
-	m.mockF(events)
-}
-
-type mockAssetFetcher struct {
-	eventsToPublish []AssetEvent
-}
-
-func (m mockAssetFetcher) Fetch(_ context.Context, assetChannel chan<- AssetEvent) {
-	for _, e := range m.eventsToPublish {
-		assetChannel <- e
-	}
-}
-
 func TestAssetInventory_Run(t *testing.T) {
 	now := func() time.Time { return time.Date(2024, 1, 1, 1, 1, 1, 0, time.Local) }
 	expected := []beat.Event{
@@ -98,58 +81,54 @@ func TestAssetInventory_Run(t *testing.T) {
 	}
 
 	publishedCh := make(chan []beat.Event)
-	publisher := mockAssetPublisher{
-		mockF: func(e []beat.Event) {
-			publishedCh <- e
-		},
-	}
+	publisher := NewMockAssetPublisher(t)
+	publisher.EXPECT().PublishAll(mock.Anything).Run(func(e []beat.Event) {
+		publishedCh <- e
+	})
 
-	fetchers := []AssetFetcher{
-		mockAssetFetcher{
-			eventsToPublish: []AssetEvent{
-				NewAssetEvent(
-					AssetClassification{
-						Category:    CategoryInfrastructure,
-						SubCategory: SubCategoryCompute,
-						Type:        TypeVirtualMachine,
-						SubStype:    SubTypeEC2,
-					},
-					"arn:aws:ec2:us-east::ec2/234567890",
-					"test-server",
-					WithTags(map[string]string{"Name": "test-server", "key": "value"}),
-					WithCloud(AssetCloud{
-						Provider: AwsCloudProvider,
-						Region:   "us-east",
-					}),
-					WithHost(AssetHost{
-						Architecture:    string(types.ArchitectureValuesX8664),
-						ImageId:         pointers.Ref("image-id"),
-						InstanceType:    "instance-type",
-						Platform:        "linux",
-						PlatformDetails: pointers.Ref("ubuntu"),
-					}),
-					WithIAM(AssetIAM{
-						Id:  pointers.Ref("a123123"),
-						Arn: pointers.Ref("123123:123123:123123"),
-					}),
-					WithNetwork(AssetNetwork{
-						NetworkId:        pointers.Ref("vpc-id"),
-						SubnetId:         pointers.Ref("subnetId"),
-						Ipv6Address:      pointers.Ref("ipv6"),
-						PublicIpAddress:  pointers.Ref("public-ip-addr"),
-						PrivateIpAddress: pointers.Ref("private-ip-addre"),
-						PublicDnsName:    pointers.Ref("public-dns"),
-						PrivateDnsName:   pointers.Ref("private-dns"),
-					}),
-				),
+	fetcher := NewMockAssetFetcher(t)
+	fetcher.EXPECT().Fetch(mock.Anything, mock.Anything).Run(func(_ context.Context, assetChannel chan<- AssetEvent) {
+		assetChannel <- NewAssetEvent(
+			AssetClassification{
+				Category:    CategoryInfrastructure,
+				SubCategory: SubCategoryCompute,
+				Type:        TypeVirtualMachine,
+				SubStype:    SubTypeEC2,
 			},
-		},
-	}
+			"arn:aws:ec2:us-east::ec2/234567890",
+			"test-server",
+			WithTags(map[string]string{"Name": "test-server", "key": "value"}),
+			WithCloud(AssetCloud{
+				Provider: AwsCloudProvider,
+				Region:   "us-east",
+			}),
+			WithHost(AssetHost{
+				Architecture:    string(types.ArchitectureValuesX8664),
+				ImageId:         pointers.Ref("image-id"),
+				InstanceType:    "instance-type",
+				Platform:        "linux",
+				PlatformDetails: pointers.Ref("ubuntu"),
+			}),
+			WithIAM(AssetIAM{
+				Id:  pointers.Ref("a123123"),
+				Arn: pointers.Ref("123123:123123:123123"),
+			}),
+			WithNetwork(AssetNetwork{
+				NetworkId:        pointers.Ref("vpc-id"),
+				SubnetId:         pointers.Ref("subnetId"),
+				Ipv6Address:      pointers.Ref("ipv6"),
+				PublicIpAddress:  pointers.Ref("public-ip-addr"),
+				PrivateIpAddress: pointers.Ref("private-ip-addre"),
+				PublicDnsName:    pointers.Ref("public-dns"),
+				PrivateDnsName:   pointers.Ref("private-dns"),
+			}),
+		)
+	})
 
 	logger := logp.NewLogger("test_run")
 	inventory := AssetInventory{
 		logger:              logger,
-		fetchers:            fetchers,
+		fetchers:            []AssetFetcher{fetcher},
 		publisher:           publisher,
 		bufferFlushInterval: 10 * time.Millisecond,
 		bufferMaxSize:       1,

diff --git a/internal/inventory/mock_asset_enricher.go b/internal/inventory/mock_asset_enricher.go