trick/types_test.go at main · wakeful/trick · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
// Copyright 2025 variHQ OÜ
// SPDX-License-Identifier: BSD-3-Clause

package main

import (
	"context"
	"reflect"
	"testing"

	"github.com/aws/aws-sdk-go-v2/service/sts"
)

type MockSTSClient struct {
	mockGetCallerIdentityOutput sts.GetCallerIdentityOutput
	mockGetCallerIdentityError  error
	mockAssumeRoleOutput        map[string]sts.AssumeRoleOutput
	mockAssumeRoleError         error
}

func (m MockSTSClient) AssumeRole(
	_ context.Context,
	params *sts.AssumeRoleInput,
	_ ...func(*sts.Options),
) (*sts.AssumeRoleOutput, error) {
	v, ok := m.mockAssumeRoleOutput[*params.RoleArn]
	if !ok {
		return nil, m.mockAssumeRoleError
	}

	return &v, nil
}

func (m MockSTSClient) GetCallerIdentity(
	_ context.Context,
	_ *sts.GetCallerIdentityInput,
	_ ...func(*sts.Options),
) (*sts.GetCallerIdentityOutput, error) {
	return &m.mockGetCallerIdentityOutput, m.mockGetCallerIdentityError
}

var _ ServiceSTS = (*MockSTSClient)(nil)

func TestNewApp(t *testing.T) {
	t.Parallel()

	pool, err := setRolePool([]string{"role-a", "role-b"})
	if err != nil {
		t.Fatalf("setRolePool failed: %v", err)
	}

	type args struct {
		region      string
		roles       []string
		usableRoles []string
	}

	tests := []struct {
		name    string
		args    args
		want    *App
		wantErr bool
	}{
		{
			name: "success",
			args: args{
				region:      "eu-west-1",
				roles:       []string{"role-a", "role-b"},
				usableRoles: []string{"role-a", "role-b"},
			},
			want: &App{
				client:          MockSTSClient{},
				region:          "eu-west-1",
				roles:           pool,
				usableRoles:     map[string]struct{}{"role-a": {}, "role-b": {}},
				sessionDuration: 0,
			},
			wantErr: false,
		},
		{
			name: "error",
			args: args{
				region:      "eu-west-1",
				roles:       []string{"role-a", "role-b"},
				usableRoles: []string{"role-c"},
			},
			want:    nil,
			wantErr: true,
		},
		{
			name: "no roles",
			args: args{
				region:      "eu-west-1",
				roles:       []string{},
				usableRoles: []string{},
			},
			want:    nil,
			wantErr: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			got, err := NewApp(t.Context(), tt.args.region, tt.args.roles, tt.args.usableRoles)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewApp() error = %v, wantErr %v", err, tt.wantErr)

				return
			}

			if (err != nil) && tt.wantErr {
				return
			}

			if !reflect.DeepEqual(got.roles, tt.want.roles) {
				t.Errorf("NewApp() roles got = %v, want %v", got, tt.want)
			}

			if !reflect.DeepEqual(got.usableRoles, tt.want.usableRoles) {
				t.Errorf("NewApp() usableRoles got = %v, want %v", got, tt.want)
			}
		})
	}
}