@@ -9,19 +9,22 @@ package ec2
9
9
10
10
import (
11
11
"context"
12
+ "errors"
12
13
"fmt"
13
14
"io"
14
15
"net/http"
15
16
"net/http/httptest"
16
17
"os"
17
18
"testing"
18
19
19
- "github.com/stretchr/testify/assert"
20
- "github.com/stretchr/testify/require"
21
-
22
20
configmock "github.com/DataDog/datadog-agent/pkg/config/mock"
23
21
pkgconfigsetup "github.com/DataDog/datadog-agent/pkg/config/setup"
24
22
"github.com/DataDog/datadog-agent/pkg/util/cache"
23
+ "github.com/aws/aws-sdk-go-v2/aws"
24
+ "github.com/aws/aws-sdk-go-v2/service/ec2"
25
+ "github.com/aws/aws-sdk-go-v2/service/ec2/types"
26
+ "github.com/stretchr/testify/assert"
27
+ "github.com/stretchr/testify/require"
25
28
)
26
29
27
30
func TestGetIAMRole (t * testing.T ) {
@@ -208,3 +211,105 @@ func TestGetTagsFullWorkflow(t *testing.T) {
208
211
assert .NoError (t , err )
209
212
assert .Equal (t , []string {"tag1" , "tag2" }, tags )
210
213
}
214
+
215
+ // Mock implementation of ec2ClientInterface
216
+ type mockEC2Client struct {
217
+ DescribeTagsFunc func (ctx context.Context , params * ec2.DescribeTagsInput , optFns ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error )
218
+ }
219
+
220
+ func (m * mockEC2Client ) DescribeTags (ctx context.Context , params * ec2.DescribeTagsInput , optFns ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error ) {
221
+ return m .DescribeTagsFunc (ctx , params , optFns ... )
222
+ }
223
+
224
+ // Helper function to compare slices of strings
225
+ func equalStringSlices (a , b []string ) bool {
226
+ if len (a ) != len (b ) {
227
+ return false
228
+ }
229
+ aMap := make (map [string ]struct {}, len (a ))
230
+ for _ , v := range a {
231
+ aMap [v ] = struct {}{}
232
+ }
233
+ for _ , v := range b {
234
+ if _ , ok := aMap [v ]; ! ok {
235
+ return false
236
+ }
237
+ }
238
+ return true
239
+ }
240
+
241
+ func TestGetTagsWithCreds (t * testing.T ) {
242
+ tests := []struct {
243
+ name string
244
+ instanceIdentity * EC2Identity
245
+ mockDescribeTags func (ctx context.Context , params * ec2.DescribeTagsInput , optFns ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error )
246
+ expectedTags []string
247
+ expectedError assert.ErrorAssertionFunc
248
+ }{
249
+ {
250
+ name : "Successful retrieval of tags" ,
251
+ instanceIdentity : & EC2Identity {
252
+ InstanceID : "i-1234567890abcdef0" ,
253
+ },
254
+ mockDescribeTags : func (_ context.Context , _ * ec2.DescribeTagsInput , _ ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error ) {
255
+ return & ec2.DescribeTagsOutput {
256
+ Tags : []types.TagDescription {
257
+ {Key : aws .String ("Name" ), Value : aws .String ("TestInstance" )},
258
+ {Key : aws .String ("Env" ), Value : aws .String ("Production" )},
259
+ },
260
+ }, nil
261
+ },
262
+ expectedTags : []string {"Name:TestInstance" , "Env:Production" },
263
+ expectedError : assert .NoError ,
264
+ },
265
+ {
266
+ name : "Excluded tags are filtered out" ,
267
+ instanceIdentity : & EC2Identity {
268
+ InstanceID : "i-1234567890abcdef0" ,
269
+ },
270
+ mockDescribeTags : func (_ context.Context , _ * ec2.DescribeTagsInput , _ ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error ) {
271
+ return & ec2.DescribeTagsOutput {
272
+ Tags : []types.TagDescription {
273
+ {Key : aws .String ("Name" ), Value : aws .String ("TestInstance" )},
274
+ {Key : aws .String ("aws:cloudformation:stack-name" ), Value : aws .String ("MyStack" )},
275
+ },
276
+ }, nil
277
+ },
278
+ expectedTags : []string {"Name:TestInstance" , "aws:cloudformation:stack-name:MyStack" },
279
+ expectedError : assert .NoError ,
280
+ },
281
+ {
282
+ name : "DescribeTags returns an error" ,
283
+ instanceIdentity : & EC2Identity {
284
+ InstanceID : "i-1234567890abcdef0" ,
285
+ },
286
+ mockDescribeTags : func (_ context.Context , _ * ec2.DescribeTagsInput , _ ... func (* ec2.Options )) (* ec2.DescribeTagsOutput , error ) {
287
+ return nil , errors .New ("DescribeTags error" )
288
+ },
289
+ expectedTags : nil ,
290
+ expectedError : assert .Error ,
291
+ },
292
+ }
293
+
294
+ for _ , tt := range tests {
295
+ t .Run (tt .name , func (t * testing.T ) {
296
+ // Mock the EC2 client
297
+ mockClient := & mockEC2Client {
298
+ DescribeTagsFunc : tt .mockDescribeTags ,
299
+ }
300
+
301
+ // Create a background context
302
+ ctx := context .Background ()
303
+
304
+ // Call the function under test
305
+ tags , err := getTagsWithCreds (ctx , tt .instanceIdentity , mockClient )
306
+
307
+ // Validate the error
308
+ tt .expectedError (t , err )
309
+ // Validate the tags
310
+ if ! equalStringSlices (tags , tt .expectedTags ) {
311
+ t .Fatalf ("Expected tags '%v', got '%v'" , tt .expectedTags , tags )
312
+ }
313
+ })
314
+ }
315
+ }
0 commit comments