refactor: fix tests

SpecterOps · Nov 1, 2022 · a6ff55b · a6ff55b
1 parent f6bb0fa
commit a6ff55b
Show file tree

Hide file tree

Showing 14 changed files with 93 additions and 137 deletions.
diff --git a/cmd/list-group-members_test.go b/cmd/list-group-members_test.go
@@ -96,6 +96,6 @@ func TestListGroupMembers(t *testing.T) {
 	} else if data, ok := wrapper.Data.(models.GroupMembers); !ok {
 		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.GroupMembers{})
 	} else if len(data.Members) != 1 {
-		t.Errorf("got %v, want %v", len(data.Members), 2)
+		t.Errorf("got %v, want %v", len(data.Members), 1)
 	}
 }
diff --git a/cmd/list-management-group-owners_test.go b/cmd/list-management-group-owners_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListManagementGroupOwners(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.ManagementGroupRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listManagementGroupOwners(ctx, mockClient, mockRoleAssignmentsChannel)
+	channel := listManagementGroupOwners(ctx, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.ManagementGroupRoleAssignments{
+		mockRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZManagementGroupRoleAssignment,
+			models.ManagementGroupRoleAssignments{
 				ManagementGroupId: "foo",
 				RoleAssignments: []models.ManagementGroupRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListManagementGroupOwners(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.ManagementGroupOwners); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ManagementGroupOwners{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-management-group-role-assignments_test.go b/cmd/list-management-group-role-assignments_test.go
@@ -93,21 +93,13 @@ func TestListResourceGroupRoleAssignments(t *testing.T) {
 
 	if result, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if data, ok := wrapper.Data.(models.ResourceGroupRoleAssignments); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ResourceGroupRoleAssignments{})
-	} else if len(data.RoleAssignments) != 2 {
-		t.Errorf("got %v, want %v", len(data.RoleAssignments), 2)
+	} else if len(result.Data.RoleAssignments) != 2 {
+		t.Errorf("got %v, want %v", len(result.Data.RoleAssignments), 2)
 	}
 
 	if result, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if data, ok := wrapper.Data.(models.ResourceGroupRoleAssignments); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ResourceGroupRoleAssignments{})
-	} else if len(data.RoleAssignments) != 1 {
-		t.Errorf("got %v, want %v", len(data.RoleAssignments), 2)
+	} else if len(result.Data.RoleAssignments) != 1 {
+		t.Errorf("got %v, want %v", len(result.Data.RoleAssignments), 2)
 	}
 }
diff --git a/cmd/list-management-group-user-access-admins_test.go b/cmd/list-management-group-user-access-admins_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListManagementGroupUserAccessAdmins(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.ManagementGroupRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listManagementGroupUserAccessAdmins(ctx, mockClient, mockRoleAssignmentsChannel)
+	channel := listManagementGroupUserAccessAdmins(ctx, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.ManagementGroupRoleAssignments{
+		mockRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZManagementGroupRoleAssignment,
+			models.ManagementGroupRoleAssignments{
 				ManagementGroupId: "foo",
 				RoleAssignments: []models.ManagementGroupRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListManagementGroupUserAccessAdmins(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.ManagementGroupUserAccessAdmins); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ManagementGroupUserAccessAdmins{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-resource-group-owners_test.go b/cmd/list-resource-group-owners_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListResourceGroupOwners(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.ResourceGroupRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listResourceGroupOwners(ctx, mockClient, mockRoleAssignmentsChannel)
+	channel := listResourceGroupOwners(ctx, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.ResourceGroupRoleAssignments{
+		mockRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZResourceGroupRoleAssignment,
+			models.ResourceGroupRoleAssignments{
 				ResourceGroupId: "foo",
 				RoleAssignments: []models.ResourceGroupRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListResourceGroupOwners(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.ResourceGroupOwners); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ResourceGroupOwners{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-resource-group-role-assignments_test.go b/cmd/list-resource-group-role-assignments_test.go
@@ -93,21 +93,13 @@ func TestListManagementGroupRoleAssignments(t *testing.T) {
 
 	if result, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if data, ok := wrapper.Data.(models.ManagementGroupRoleAssignments); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ManagementGroupRoleAssignments{})
-	} else if len(data.RoleAssignments) != 2 {
-		t.Errorf("got %v, want %v", len(data.RoleAssignments), 2)
+	} else if len(result.Data.RoleAssignments) != 2 {
+		t.Errorf("got %v, want %v", len(result.Data.RoleAssignments), 2)
 	}
 
 	if result, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if data, ok := wrapper.Data.(models.ManagementGroupRoleAssignments); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ManagementGroupRoleAssignments{})
-	} else if len(data.RoleAssignments) != 1 {
-		t.Errorf("got %v, want %v", len(data.RoleAssignments), 2)
+	} else if len(result.Data.RoleAssignments) != 1 {
+		t.Errorf("got %v, want %v", len(result.Data.RoleAssignments), 2)
 	}
 }
diff --git a/cmd/list-resource-group-user-access-admins_test.go b/cmd/list-resource-group-user-access-admins_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListResourceGroupUserAccessAdmins(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockRoleAssignmentsChannel := make(chan interface{})
+	mockRoleAssignmentsChannel := make(chan azureWrapper[models.ResourceGroupRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listResourceGroupUserAccessAdmins(ctx, mockClient, mockRoleAssignmentsChannel)
+	channel := listResourceGroupUserAccessAdmins(ctx, mockRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockRoleAssignmentsChannel)
 
-		mockRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.ResourceGroupRoleAssignments{
+		mockRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZResourceGroupRoleAssignment,
+			models.ResourceGroupRoleAssignments{
 				ResourceGroupId: "foo",
 				RoleAssignments: []models.ResourceGroupRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListResourceGroupUserAccessAdmins(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.ResourceGroupUserAccessAdmins); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.ResourceGroupUserAccessAdmins{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-virtual-machine-admin-logins_test.go b/cmd/list-virtual-machine-admin-logins_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListVirtualMachineAdminLogins(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockVMRoleAssignmentsChannel := make(chan interface{})
+	mockVMRoleAssignmentsChannel := make(chan azureWrapper[models.VirtualMachineRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listVirtualMachineAdminLogins(ctx, mockClient, mockVMRoleAssignmentsChannel)
+	channel := listVirtualMachineAdminLogins(ctx, mockVMRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockVMRoleAssignmentsChannel)
 
-		mockVMRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.VirtualMachineRoleAssignments{
+		mockVMRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZVMRoleAssignment,
+			models.VirtualMachineRoleAssignments{
 				VirtualMachineId: "foo",
 				RoleAssignments: []models.VirtualMachineRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListVirtualMachineAdminLogins(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.VirtualMachineAdminLogins); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.VirtualMachineAdminLogins{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-virtual-machine-avere-contributors_test.go b/cmd/list-virtual-machine-avere-contributors_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListVirtualMachineAvereContributors(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockVMRoleAssignmentsChannel := make(chan interface{})
+	mockVMRoleAssignmentsChannel := make(chan azureWrapper[models.VirtualMachineRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listVirtualMachineAvereContributors(ctx, mockClient, mockVMRoleAssignmentsChannel)
+	channel := listVirtualMachineAvereContributors(ctx, mockVMRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockVMRoleAssignmentsChannel)
 
-		mockVMRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.VirtualMachineRoleAssignments{
+		mockVMRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZVMRoleAssignment,
+			models.VirtualMachineRoleAssignments{
 				VirtualMachineId: "foo",
 				RoleAssignments: []models.VirtualMachineRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListVirtualMachineAvereContributors(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.VirtualMachineAvereContributors); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.VirtualMachineAvereContributors{})
 	}
 
 	if _, ok := <-channel; ok {

diff --git a/cmd/list-virtual-machine-contributors_test.go b/cmd/list-virtual-machine-contributors_test.go
@@ -23,6 +23,7 @@ import (
 
 	"github.com/bloodhoundad/azurehound/client/mocks"
 	"github.com/bloodhoundad/azurehound/constants"
+	"github.com/bloodhoundad/azurehound/enums"
 	"github.com/bloodhoundad/azurehound/models"
 	"github.com/bloodhoundad/azurehound/models/azure"
 	"github.com/golang/mock/gomock"
@@ -39,16 +40,17 @@ func TestListVirtualMachineContributors(t *testing.T) {
 
 	mockClient := mocks.NewMockAzureClient(ctrl)
 
-	mockVMRoleAssignmentsChannel := make(chan interface{})
+	mockVMRoleAssignmentsChannel := make(chan azureWrapper[models.VirtualMachineRoleAssignments])
 	mockTenant := azure.Tenant{}
 	mockClient.EXPECT().TenantInfo().Return(mockTenant).AnyTimes()
-	channel := listVirtualMachineContributors(ctx, mockClient, mockVMRoleAssignmentsChannel)
+	channel := listVirtualMachineContributors(ctx, mockVMRoleAssignmentsChannel)
 
 	go func() {
 		defer close(mockVMRoleAssignmentsChannel)
 
-		mockVMRoleAssignmentsChannel <- AzureWrapper{
-			Data: models.VirtualMachineRoleAssignments{
+		mockVMRoleAssignmentsChannel <- NewAzureWrapper(
+			enums.KindAZVMRoleAssignment,
+			models.VirtualMachineRoleAssignments{
 				VirtualMachineId: "foo",
 				RoleAssignments: []models.VirtualMachineRoleAssignment{
 					{
@@ -61,15 +63,11 @@ func TestListVirtualMachineContributors(t *testing.T) {
 					},
 				},
 			},
-		}
+		)
 	}()
 
-	if result, ok := <-channel; !ok {
+	if _, ok := <-channel; !ok {
 		t.Fatalf("failed to receive from channel")
-	} else if wrapper, ok := result.(AzureWrapper); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", result, AzureWrapper{})
-	} else if _, ok := wrapper.Data.(models.VirtualMachineContributors); !ok {
-		t.Errorf("failed type assertion: got %T, want %T", wrapper.Data, models.VirtualMachineContributors{})
 	}
 
 	if _, ok := <-channel; ok {