From 819bdc13b88140abc2f52301d3d3a14e34db1350 Mon Sep 17 00:00:00 2001 From: Amber Brown Date: Thu, 26 Mar 2026 13:36:28 +1100 Subject: [PATCH 1/5] azureclient/resourceskus: replace a big ole list accumulation with an iterator --- .../armcompute/resourceskus_addons.go | 25 ++++++++++++------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/pkg/util/azureclient/azuresdk/armcompute/resourceskus_addons.go b/pkg/util/azureclient/azuresdk/armcompute/resourceskus_addons.go index e78bec6d24d..e9688f28535 100644 --- a/pkg/util/azureclient/azuresdk/armcompute/resourceskus_addons.go +++ b/pkg/util/azureclient/azuresdk/armcompute/resourceskus_addons.go @@ -5,6 +5,7 @@ package armcompute import ( "context" + "iter" armcompute "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" @@ -12,10 +13,10 @@ import ( ) type ResourceSKUsClientAddons interface { - List(ctx context.Context, filter string, includeExtendedLocations bool) ([]*armcompute.ResourceSKU, error) + List(ctx context.Context, filter string, includeExtendedLocations bool) iter.Seq2[*armcompute.ResourceSKU, error] } -func (c *resourceSKUsClient) List(ctx context.Context, filter string, includeExtendedLocations bool) (result []*armcompute.ResourceSKU, err error) { +func (c *resourceSKUsClient) List(ctx context.Context, filter string, includeExtendedLocations bool) iter.Seq2[*armcompute.ResourceSKU, error] { ex := "false" if includeExtendedLocations { ex = "true" @@ -26,13 +27,19 @@ func (c *resourceSKUsClient) List(ctx context.Context, filter string, includeExt IncludeExtendedLocations: pointerutils.ToPtr(ex), }) - for pager.More() { - page, err := pager.NextPage(ctx) - if err != nil { - return nil, err + return func(yield func(*armcompute.ResourceSKU, error) bool) { + for pager.More() { + page, err := pager.NextPage(ctx) + if err != nil { + yield(nil, err) + return + } + + for _, v := range page.Value { + if !yield(v, nil) { + return + } + } } - result = append(result, page.Value...) } - - return result, nil } From ff1071c51c3441029018bcbb0cc05ad0fe0575cd Mon Sep 17 00:00:00 2001 From: Amber Brown Date: Thu, 26 Mar 2026 14:46:48 +1100 Subject: [PATCH 2/5] replace computeskus.GetVMSkusForCurrentRegion with something that uses the iterators more efficiently --- pkg/util/computeskus/computeskus.go | 92 ++-- pkg/util/computeskus/computeskus_test.go | 559 +++++++++++++++++++---- 2 files changed, 523 insertions(+), 128 deletions(-) diff --git a/pkg/util/computeskus/computeskus.go b/pkg/util/computeskus/computeskus.go index 683e72c69c8..50817338982 100644 --- a/pkg/util/computeskus/computeskus.go +++ b/pkg/util/computeskus/computeskus.go @@ -7,7 +7,7 @@ import ( "context" "errors" "fmt" - "strings" + "slices" sdkcompute "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" @@ -65,50 +65,76 @@ func IsRestricted(sku *sdkcompute.ResourceSKU, location string) bool { return false } -// FilterVMSizes filters resource SKU by location and returns only virtual machines, their names, restrictions, location info, and capabilities. -func FilterVMSizes(skus []*sdkcompute.ResourceSKU, location string) map[string]*sdkcompute.ResourceSKU { +// SupportedOSDisk returns the type of OSDisk for the given resource. Most VMs will use Premium disks but some SKUs only support Standard SSDs +func SupportedOSDisk(vmSku *sdkcompute.ResourceSKU) string { + if HasCapability(vmSku, premiumDiskCapability) { + return premiumDisk + } + return standardDisk +} + +func SelectVMSkusInCurrentRegion(ctx context.Context, resourceSkusClient armcompute.ResourceSKUsClient, location string, skuNames []string) (map[string]*sdkcompute.ResourceSKU, error) { + // Sort and compact so that we only have one instance of each SKU in the list + slices.Sort(skuNames) + skuNames = slices.Compact(skuNames) + vmskus := map[string]*sdkcompute.ResourceSKU{} - for _, sku := range skus { - // TODO(mjudeikis): At some point some SKU's stopped returning zones and - // locations. IcM is open with MSFT but this might take a while. - // Revert once we find out right behaviour. - // https://github.com/Azure/ARO-RP/issues/1515 - if len(sku.Locations) == 0 || !strings.EqualFold(*(sku.Locations)[0], location) || - *sku.ResourceType != "virtualMachines" { - continue + filter := fmt.Sprintf("location eq %s", location) + skusIter := resourceSkusClient.List(ctx, filter, false) + + for sku, err := range skusIter { + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrListVMResourceSKUs, err) } - if len(sku.LocationInfo) == 0 { // happened in eastus2euap + // We only care about VMs and ones with locations/locationinfo + if *sku.ResourceType != "virtualMachines" || len(sku.Locations) == 0 || len(sku.LocationInfo) == 0 { continue } - // We copy only part of the object so we don't have to keep - // a lot of data in memory. - vmskus[*sku.Name] = &sdkcompute.ResourceSKU{ - Name: sku.Name, - Restrictions: sku.Restrictions, - LocationInfo: sku.LocationInfo, - Capabilities: sku.Capabilities, - } - } + if slices.Contains(skuNames, *sku.Name) { + // Make sure it's actually in our location + if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return *s == location }) { + continue + } - return vmskus -} + vmskus[*sku.Name] = sku + } -// SupportedOSDisk returns the type of OSDisk for the given resource. Most VMs will use Premium disks but some SKUs only support Standard SSDs -func SupportedOSDisk(vmSku *sdkcompute.ResourceSKU) string { - if HasCapability(vmSku, premiumDiskCapability) { - return premiumDisk + // If we've already found all the SKUs we want, exit + if len(vmskus) == len(skuNames) { + break + } } - return standardDisk + + return vmskus, nil } -func GetVMSkusForCurrentRegion(ctx context.Context, resourceSkusClient armcompute.ResourceSKUsClient, location string) (map[string]*sdkcompute.ResourceSKU, error) { +func ListUnrestrictedVMSkusInCurrentRegion(ctx context.Context, resourceSkusClient armcompute.ResourceSKUsClient, location string) ([]string, error) { + vmskus := []string{} filter := fmt.Sprintf("location eq %s", location) - skus, err := resourceSkusClient.List(ctx, filter, false) - if err != nil { - return nil, errors.Join(ErrListVMResourceSKUs, err) + skusIter := resourceSkusClient.List(ctx, filter, false) + + for sku, err := range skusIter { + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrListVMResourceSKUs, err) + } + + // We only care about VMs and ones with locations/locationinfo + if *sku.ResourceType != "virtualMachines" || len(sku.Locations) == 0 || len(sku.LocationInfo) == 0 { + continue + } + + // Make sure it's actually in our location + if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return *s == location }) { + continue + } + + if !IsRestricted(sku, location) { + vmskus = append(vmskus, *sku.Name) + } } - return FilterVMSizes(skus, location), nil + slices.Sort(vmskus) + return slices.Compact(vmskus), nil } diff --git a/pkg/util/computeskus/computeskus_test.go b/pkg/util/computeskus/computeskus_test.go index 09b81bbf5b2..bc70228965e 100644 --- a/pkg/util/computeskus/computeskus_test.go +++ b/pkg/util/computeskus/computeskus_test.go @@ -4,12 +4,18 @@ package computeskus // Licensed under the Apache License 2.0. import ( + "errors" + "maps" "reflect" "testing" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" "github.com/Azure/ARO-RP/pkg/util/cmp" + mock_armcompute "github.com/Azure/ARO-RP/pkg/util/mocks/azureclient/azuresdk/armcompute" "github.com/Azure/ARO-RP/pkg/util/pointerutils" ) @@ -103,101 +109,6 @@ func TestHasCapability(t *testing.T) { } } -func TestFilterVmSizes(t *testing.T) { - for _, tt := range []struct { - name string - providedLocation string - resourceType string - skuLocation []string - skuRestrictions armcompute.ResourceSKURestrictions - skuLocationInfo []armcompute.ResourceSKULocationInfo - skuCapabilities string - wantResult map[string]*armcompute.ResourceSKU - }{ - { - name: "resource type is a virtual machine", - providedLocation: "eastus", - resourceType: "virtualMachines", - skuRestrictions: armcompute.ResourceSKURestrictions{ReasonCode: pointerutils.ToPtr(armcompute.ResourceSKURestrictionsReasonCodeNotAvailableForSubscription)}, - skuLocation: []string{"eastus"}, - skuLocationInfo: []armcompute.ResourceSKULocationInfo{{Zones: pointerutils.ToSlicePtr([]string{"eastus-2"})}}, - skuCapabilities: "some-capability", - - wantResult: map[string]*armcompute.ResourceSKU{ - "Fake_Sku": { - Name: pointerutils.ToPtr("Fake_Sku"), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{{ - ReasonCode: pointerutils.ToPtr(armcompute.ResourceSKURestrictionsReasonCodeNotAvailableForSubscription), - }}), - LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{ - { - Zones: pointerutils.ToSlicePtr([]string{"eastus-2"}), - }, - }), - Capabilities: pointerutils.ToSlicePtr([]armcompute.ResourceSKUCapabilities{{ - Name: pointerutils.ToPtr("some-capability"), - }}), - }, - }, - }, - { - name: "resource type not a virtual machine", - providedLocation: "eastus", - resourceType: "disk", - skuLocation: []string{"eastus"}, - skuLocationInfo: []armcompute.ResourceSKULocationInfo{{Zones: pointerutils.ToSlicePtr([]string{"eastus-2"})}}, - wantResult: map[string]*armcompute.ResourceSKU{}, - }, - { - name: "sku Location doesn't match provided location", - providedLocation: "mars", - resourceType: "virtualMachines", - skuLocation: []string{"eastus"}, - skuLocationInfo: []armcompute.ResourceSKULocationInfo{{Zones: pointerutils.ToSlicePtr([]string{"eastus-2"})}}, - wantResult: map[string]*armcompute.ResourceSKU{}, - }, - { - name: "sku Location has length of 0", - providedLocation: "eastus", - resourceType: "virtualMachines", - skuLocation: []string{}, - skuLocationInfo: []armcompute.ResourceSKULocationInfo{{Zones: pointerutils.ToSlicePtr([]string{"eastus-2"})}}, - wantResult: map[string]*armcompute.ResourceSKU{}, - }, - { - name: "sku LocationInfo has length of 0", - providedLocation: "eastus", - resourceType: "virtualMachines", - skuLocation: []string{"eastus"}, - skuLocationInfo: []armcompute.ResourceSKULocationInfo{}, - wantResult: map[string]*armcompute.ResourceSKU{}, - }, - } { - t.Run(tt.name, func(t *testing.T) { - sku := []*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr("Fake_Sku"), - Capabilities: pointerutils.ToSlicePtr([]armcompute.ResourceSKUCapabilities{ - { - Name: pointerutils.ToPtr(tt.skuCapabilities), - }, - }), - Locations: pointerutils.ToSlicePtr(tt.skuLocation), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{tt.skuRestrictions}), - LocationInfo: pointerutils.ToSlicePtr(tt.skuLocationInfo), - ResourceType: pointerutils.ToPtr(tt.resourceType), - }, - } - - result := FilterVMSizes(sku, tt.providedLocation) - - if !reflect.DeepEqual(result, tt.wantResult) { - t.Error(cmp.Diff(result, tt.wantResult)) - } - }) - } -} - func TestIsRestricted(t *testing.T) { for _, tt := range []struct { name string @@ -297,3 +208,461 @@ func TestSupportedOSDisk(t *testing.T) { }) } } + +func TestSelectVMSkusInCurrentRegion(t *testing.T) { + for _, tt := range []struct { + name string + vmSkus []string + mocks func(*mock_armcompute.MockResourceSKUsClient) + desired map[string]*armcompute.ResourceSKU + wantError string + }{ + { + name: "happypath", + vmSkus: []string{"bigmachine_v1", "smallmachine_v4", "smallmachine_v5"}, + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + // Not actually in our region + { + Name: pointerutils.ToPtr("smallmachine_v5"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus1"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus1"), + }, + }, + }: nil, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }: nil, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + }), + ) + }, + desired: map[string]*armcompute.ResourceSKU{ + "bigmachine_v1": { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + "smallmachine_v4": { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + }, + }, + { + name: "duplicate skus in vmskus", + vmSkus: []string{"bigmachine_v1", "bigmachine_v1", "smallmachine_v4"}, + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{}, + }: nil, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }: nil, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + }), + ) + }, + desired: map[string]*armcompute.ResourceSKU{ + "bigmachine_v1": { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + "smallmachine_v4": { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + }, + }, + { + name: "error in pagination", + vmSkus: []string{"bigmachine_v5"}, + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + nil: errors.New("this is an error"), + }), + ) + }, + wantError: "this is an error", + }, + } { + t.Run(tt.name, func(t *testing.T) { + controller := gomock.NewController(t) + mock_skus := mock_armcompute.NewMockResourceSKUsClient(controller) + tt.mocks(mock_skus) + + r, err := SelectVMSkusInCurrentRegion(t.Context(), mock_skus, "northus2", tt.vmSkus) + if tt.wantError != "" { + require.ErrorContains(t, err, tt.wantError) + } else { + require.NoError(t, err) + require.Equal(t, tt.desired, r) + } + }) + } +} + +func TestListUnrestrictedSKUNames(t *testing.T) { + for _, tt := range []struct { + name string + mocks func(*mock_armcompute.MockResourceSKUsClient) + desired []string + wantError string + }{ + { + name: "happypath", + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + // Duplicated struct, in case we get two + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{}, + }: nil, + // Not actually in our region + { + Name: pointerutils.ToPtr("smallmachine_v9"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus1"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus1"), + }, + }, + }: nil, + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, + }, + }, + }: nil, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }: nil, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + }), + ) + }, + desired: []string{"bigmachine_v1", "smallmachine_v4"}, + }, + { + name: "duplicate VM structs don't lead to duplicated names", + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + {}, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, + }, + }, + }: nil, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }: nil, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + }), + ) + }, + desired: []string{"bigmachine_v1", "smallmachine_v4"}, + }, + { + name: "error in pagination", + mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( + maps.All(map[*armcompute.ResourceSKU]error{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, + }, + }, + }: nil, + nil: errors.New("this is an error"), + }), + ) + }, + wantError: "this is an error", + }, + } { + t.Run(tt.name, func(t *testing.T) { + controller := gomock.NewController(t) + mock_skus := mock_armcompute.NewMockResourceSKUsClient(controller) + tt.mocks(mock_skus) + + r, err := ListUnrestrictedVMSkusInCurrentRegion(t.Context(), mock_skus, "northus2") + if tt.wantError != "" { + require.ErrorContains(t, err, tt.wantError) + } else { + require.NoError(t, err) + require.Equal(t, tt.desired, r) + } + }) + } +} From 95a48557fc709b03f473cc42d610c1407ac8aec1 Mon Sep 17 00:00:00 2001 From: Amber Brown Date: Thu, 26 Mar 2026 15:04:44 +1100 Subject: [PATCH 3/5] update uses of the computeskus to use the new APIs --- pkg/cluster/loadbalancerinternal.go | 2 +- pkg/cluster/loadbalancerinternal_test.go | 66 +++++++++---------- pkg/cluster/validate.go | 6 +- pkg/cluster/validate_test.go | 15 ++++- ...penshiftcluster_vmresize_pre_validation.go | 6 +- ...iftcluster_vmresize_pre_validation_test.go | 29 +++----- .../admin_openshiftcluster_vmsizelist.go | 20 ++---- .../admin_openshiftcluster_vmsizelist_test.go | 20 +----- pkg/frontend/adminactions/azureactions.go | 13 ++-- pkg/frontend/sku_test.go | 14 +++- pkg/frontend/sku_validation.go | 14 ++-- pkg/util/mocks/adminactions/azureactions.go | 19 +++++- .../azuresdk/armcompute/armcompute.go | 8 +-- 13 files changed, 115 insertions(+), 117 deletions(-) diff --git a/pkg/cluster/loadbalancerinternal.go b/pkg/cluster/loadbalancerinternal.go index f1314e18d08..7134d58639f 100644 --- a/pkg/cluster/loadbalancerinternal.go +++ b/pkg/cluster/loadbalancerinternal.go @@ -60,7 +60,7 @@ func MigrateInternalLoadBalancerZones( } } - filteredSkus, err := computeskus.GetVMSkusForCurrentRegion(ctx, resourceSkusClient, location) + filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, resourceSkusClient, location, []string{string(doc.OpenShiftCluster.Properties.MasterProfile.VMSize)}) if err != nil { return doc, err } diff --git a/pkg/cluster/loadbalancerinternal_test.go b/pkg/cluster/loadbalancerinternal_test.go index d03b87cf80c..0e221db963d 100644 --- a/pkg/cluster/loadbalancerinternal_test.go +++ b/pkg/cluster/loadbalancerinternal_test.go @@ -94,19 +94,20 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) { }, nil, ) - sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)), - Locations: pointerutils.ToSlicePtr([]string{"eastus"}), - LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{ - { - Zones: []*string{}, - }, - }), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), - ResourceType: pointerutils.ToPtr("virtualMachines"), - }, - }, nil) + sku.EXPECT().List(gomock.Any(), "location eq eastus", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + yield(&armcompute.ResourceSKU{ + Name: pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)), + Locations: pointerutils.ToSlicePtr([]string{"eastus"}), + LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{ + { + Zones: []*string{}, + }, + }), + Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), + ResourceType: pointerutils.ToPtr("virtualMachines"), + }, nil) + }) }, expectedLogs: []testlog.ExpectedLogEntry{ { @@ -136,15 +137,8 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) { }, nil, ) - sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)), - Locations: pointerutils.ToSlicePtr([]string{"eastus"}), - LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{}), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), - ResourceType: pointerutils.ToPtr("virtualMachines"), - }, - }, nil) + sku.EXPECT().List(gomock.Any(), "location eq eastus", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) {}) }, expectedLogs: []testlog.ExpectedLogEntry{}, wantErrs: []error{errVMAvailability}, @@ -170,7 +164,10 @@ func TestUpdateLoadBalancerZonalNoopAndErrorPaths(t *testing.T) { }, nil, ) - sku.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{}, errTestSKUFetchError) + sku.EXPECT().List(gomock.Any(), "location eq eastus", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + yield(nil, errTestSKUFetchError) + }) }, expectedLogs: []testlog.ExpectedLogEntry{}, wantErrs: []error{computeskus.ErrListVMResourceSKUs, errTestSKUFetchError}, @@ -380,17 +377,18 @@ func TestUpdateLoadBalancerZonalMigration(t *testing.T) { }, nil, ) - skus.EXPECT().List(gomock.Any(), "location eq eastus", false).Return([]*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)), - Locations: pointerutils.ToSlicePtr([]string{"eastus"}), - LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{ - {Zones: pointerutils.ToSlicePtr([]string{"1", "2", "3"})}, - }), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), - ResourceType: pointerutils.ToPtr("virtualMachines"), - }, - }, nil) + skus.EXPECT().List(gomock.Any(), "location eq eastus", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + yield(&armcompute.ResourceSKU{ + Name: pointerutils.ToPtr(string(api.VMSizeStandardD16asV4)), + Locations: pointerutils.ToSlicePtr([]string{"eastus"}), + LocationInfo: pointerutils.ToSlicePtr([]armcompute.ResourceSKULocationInfo{ + {Zones: pointerutils.ToSlicePtr([]string{"1", "2", "3"})}, + }), + Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), + ResourceType: pointerutils.ToPtr("virtualMachines"), + }, nil) + }) plsFIPRemoval := plses.EXPECT().CreateOrUpdateAndWait(gomock.Any(), rgName, infraID+"-pls", armnetwork.PrivateLinkService{ Properties: &armnetwork.PrivateLinkServiceProperties{ diff --git a/pkg/cluster/validate.go b/pkg/cluster/validate.go index b19f8a59ae4..f0604797e3d 100644 --- a/pkg/cluster/validate.go +++ b/pkg/cluster/validate.go @@ -32,7 +32,11 @@ func (m *manager) validateResources(ctx context.Context) error { // created. This function is only to be called during cluster bootstrap! func (m *manager) validateZones(ctx context.Context) error { location := m.doc.OpenShiftCluster.Location - filteredSkus, err := computeskus.GetVMSkusForCurrentRegion(ctx, m.armResourceSKUs, location) + + filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, m.armResourceSKUs, location, []string{ + string(m.doc.OpenShiftCluster.Properties.MasterProfile.VMSize), + string(m.doc.OpenShiftCluster.Properties.WorkerProfiles[0].VMSize), + }) if err != nil { return err } diff --git a/pkg/cluster/validate_test.go b/pkg/cluster/validate_test.go index a750c960947..accd6b480a3 100644 --- a/pkg/cluster/validate_test.go +++ b/pkg/cluster/validate_test.go @@ -6,7 +6,6 @@ package cluster import ( "context" "errors" - "fmt" "strings" "testing" @@ -180,8 +179,18 @@ func TestValidateZones(t *testing.T) { resourceSkusClient := mock_armcompute.NewMockResourceSKUsClient(controller) resourceSkusClient.EXPECT(). - List(gomock.Any(), fmt.Sprintf("location eq %v", "eastus"), false). - Return(skus, tt.resourceSkusClientErr) + List(gomock.Any(), "location eq eastus", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + if tt.resourceSkusClientErr != nil { + yield(nil, tt.resourceSkusClientErr) + return + } + for _, v := range skus { + if !yield(v, nil) { + return + } + } + }) m := &manager{ env: mockEnv, diff --git a/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation.go b/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation.go index ed3257689e4..56e1c5b043f 100644 --- a/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation.go +++ b/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation.go @@ -28,7 +28,6 @@ import ( arov1alpha1 "github.com/Azure/ARO-RP/pkg/operator/apis/aro.openshift.io/v1alpha1" "github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/compute" "github.com/Azure/ARO-RP/pkg/util/clusteroperators" - "github.com/Azure/ARO-RP/pkg/util/computeskus" ) // getPreResizeControlPlaneVMsValidation is the HTTP handler; the underscore @@ -337,15 +336,12 @@ func (f *frontend) validateVMSKU( return err } - skus, err := a.VMSizeList(ctx) + filteredSkus, err := a.VMGetSKUs(ctx, []string{desiredVMSize}) if err != nil { return err } location := doc.OpenShiftCluster.Location - - filteredSkus := computeskus.FilterVMSizes(skus, location) - sku, err := checkSKUAvailability(filteredSkus, location, "vmSize", desiredVMSize) if err != nil { return err diff --git a/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation_test.go b/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation_test.go index 0b4f08fad8e..f18175ceba7 100644 --- a/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation_test.go +++ b/pkg/frontend/admin_openshiftcluster_vmresize_pre_validation_test.go @@ -142,9 +142,9 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) { }, mocks: func(tt *test, a *mock_adminactions.MockAzureActions) { a.EXPECT(). - VMSizeList(gomock.Any()). - Return([]*armcompute.ResourceSKU{ - { + VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}). + Return(map[string]*armcompute.ResourceSKU{ + "Standard_D8s_v3": { Name: pointerutils.ToPtr("Standard_D8s_v3"), ResourceType: pointerutils.ToPtr("virtualMachines"), Locations: pointerutils.ToSlicePtr([]string{"eastus"}), @@ -308,21 +308,8 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) { }, mocks: func(tt *test, a *mock_adminactions.MockAzureActions) { a.EXPECT(). - VMSizeList(gomock.Any()). - Return([]*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr("Standard_D16s_v3"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"eastus"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("eastus"), - }, - }, - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), - Capabilities: []*armcompute.ResourceSKUCapabilities{}, - }, - }, nil) + VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}). + Return(map[string]*armcompute.ResourceSKU{}, nil) }, kubeMocks: allKubeChecksHealthyMock, wantStatusCode: http.StatusBadRequest, @@ -360,9 +347,9 @@ func TestPreResizeControlPlaneVMsValidation(t *testing.T) { }, mocks: func(tt *test, a *mock_adminactions.MockAzureActions) { a.EXPECT(). - VMSizeList(gomock.Any()). - Return([]*armcompute.ResourceSKU{ - { + VMGetSKUs(gomock.Any(), []string{"Standard_D8s_v3"}). + Return(map[string]*armcompute.ResourceSKU{ + "Standard_D8s_v3": { Name: pointerutils.ToPtr("Standard_D8s_v3"), ResourceType: pointerutils.ToPtr("virtualMachines"), Locations: pointerutils.ToSlicePtr([]string{"eastus"}), diff --git a/pkg/frontend/admin_openshiftcluster_vmsizelist.go b/pkg/frontend/admin_openshiftcluster_vmsizelist.go index 2e1bfb680d0..1bf0147bdae 100644 --- a/pkg/frontend/admin_openshiftcluster_vmsizelist.go +++ b/pkg/frontend/admin_openshiftcluster_vmsizelist.go @@ -9,13 +9,12 @@ import ( "fmt" "net/http" "path/filepath" + "slices" "strings" "github.com/go-chi/chi/v5" "github.com/sirupsen/logrus" - "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" - "github.com/Azure/ARO-RP/pkg/api" "github.com/Azure/ARO-RP/pkg/database/cosmosdb" "github.com/Azure/ARO-RP/pkg/frontend/middleware" @@ -66,19 +65,8 @@ func (f *frontend) _getAdminOpenShiftClusterVMResizeOptions(ctx context.Context, return nil, err } - return json.Marshal(f.filterVMSkus(skus)) -} - -func (f *frontend) filterVMSkus(skus []*armcompute.ResourceSKU) []string { - filteredSkus := []string{} - - for _, sku := range skus { - if sku.Restrictions != nil && len(sku.Restrictions) == 0 { - if sku.Name != nil { - filteredSkus = append(filteredSkus, *sku.Name) - } - } - } + // Sort for stability + slices.Sort(skus) - return filteredSkus + return json.Marshal(skus) } diff --git a/pkg/frontend/admin_openshiftcluster_vmsizelist_test.go b/pkg/frontend/admin_openshiftcluster_vmsizelist_test.go index 646795ab7b3..619d047ff85 100644 --- a/pkg/frontend/admin_openshiftcluster_vmsizelist_test.go +++ b/pkg/frontend/admin_openshiftcluster_vmsizelist_test.go @@ -13,14 +13,11 @@ import ( "github.com/sirupsen/logrus" "go.uber.org/mock/gomock" - "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" - "github.com/Azure/ARO-RP/pkg/api" "github.com/Azure/ARO-RP/pkg/env" "github.com/Azure/ARO-RP/pkg/frontend/adminactions" "github.com/Azure/ARO-RP/pkg/metrics/noop" mock_adminactions "github.com/Azure/ARO-RP/pkg/util/mocks/adminactions" - "github.com/Azure/ARO-RP/pkg/util/pointerutils" testdatabase "github.com/Azure/ARO-RP/test/database" ) @@ -68,23 +65,10 @@ func TestAdminListVMSizeList(t *testing.T) { mocks: func(tt *test, a *mock_adminactions.MockAzureActions) { a.EXPECT(). VMSizeList(gomock.Any()). - Return([]*armcompute.ResourceSKU{ - { - Name: pointerutils.ToPtr("Standard_D8s_v90"), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{ - { - Type: pointerutils.ToPtr(armcompute.ResourceSKURestrictionsTypeLocation), - }, - }), - }, - { - Name: pointerutils.ToPtr("Standard_D8s_v9001"), - Restrictions: pointerutils.ToSlicePtr([]armcompute.ResourceSKURestrictions{}), - }, - }, nil) + Return([]string{"Standard_D8s_v9001", "Standard_D8s_v80"}, nil) }, wantStatusCode: http.StatusOK, - wantResponse: []byte(`["Standard_D8s_v9001"]` + "\n"), + wantResponse: []byte(`["Standard_D8s_v80","Standard_D8s_v9001"]` + "\n"), }, { name: "cluster not found", diff --git a/pkg/frontend/adminactions/azureactions.go b/pkg/frontend/adminactions/azureactions.go index a9a70fd3eee..131eb635e3d 100644 --- a/pkg/frontend/adminactions/azureactions.go +++ b/pkg/frontend/adminactions/azureactions.go @@ -22,6 +22,7 @@ import ( "github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/compute" "github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/features" "github.com/Azure/ARO-RP/pkg/util/azureclient/mgmt/storage" + "github.com/Azure/ARO-RP/pkg/util/computeskus" "github.com/Azure/ARO-RP/pkg/util/stringutils" ) @@ -34,7 +35,8 @@ type AzureActions interface { VMRedeployAndWait(ctx context.Context, vmName string) error VMStartAndWait(ctx context.Context, vmName string) error VMStopAndWait(ctx context.Context, vmName string, deallocateVM bool) error - VMSizeList(ctx context.Context) ([]*sdkcompute.ResourceSKU, error) + VMSizeList(ctx context.Context) ([]string, error) + VMGetSKUs(ctx context.Context, vmSizes []string) (map[string]*sdkcompute.ResourceSKU, error) VMResize(ctx context.Context, vmName string, vmSize string) error ResourceGroupHasVM(ctx context.Context, vmName string) (bool, error) VMSerialConsole(ctx context.Context, log *logrus.Entry, vmName string, target io.Writer) error @@ -148,9 +150,12 @@ func (a *azureActions) VMStopAndWait(ctx context.Context, vmName string, dealloc return a.virtualMachines.StopAndWait(ctx, clusterRGName, vmName, deallocateVM) } -func (a *azureActions) VMSizeList(ctx context.Context) ([]*sdkcompute.ResourceSKU, error) { - filter := fmt.Sprintf("location eq '%s'", a.env.Location()) - return a.resourceSkus.List(ctx, filter, false) +func (a *azureActions) VMGetSKUs(ctx context.Context, vmSizes []string) (map[string]*sdkcompute.ResourceSKU, error) { + return computeskus.SelectVMSkusInCurrentRegion(ctx, a.resourceSkus, a.env.Location(), vmSizes) +} + +func (a *azureActions) VMSizeList(ctx context.Context) ([]string, error) { + return computeskus.ListUnrestrictedVMSkusInCurrentRegion(ctx, a.resourceSkus, a.env.Location()) } func (a *azureActions) VMResize(ctx context.Context, vmName string, size string) error { diff --git a/pkg/frontend/sku_test.go b/pkg/frontend/sku_test.go index 7c7c7953712..6e452b6438f 100644 --- a/pkg/frontend/sku_test.go +++ b/pkg/frontend/sku_test.go @@ -95,7 +95,7 @@ func TestValidateVMSku(t *testing.T) { workerProfile1Sku: "Standard_D4s_v2", workerProfile2Sku: "Standard_D4s_v2", resourceSkusClientErr: errors.New("unable to retrieve skus information"), - wantErr: "unable to retrieve skus information", + wantErr: "failure listing resource SKUs: unable to retrieve skus information", masterEncryptionAtHost: api.EncryptionAtHostDisabled, workerEncryptionAtHost: api.EncryptionAtHostDisabled, }, @@ -320,7 +320,17 @@ func TestValidateVMSku(t *testing.T) { resourceSkusClient := mock_armcompute.NewMockResourceSKUsClient(controller) resourceSkusClient.EXPECT(). List(gomock.Any(), fmt.Sprintf("location eq %v", "eastus"), false). - Return(skus, tt.resourceSkusClientErr) + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + if tt.resourceSkusClientErr != nil { + yield(nil, tt.resourceSkusClientErr) + return + } + for _, v := range skus { + if !yield(v, nil) { + return + } + } + }) err := validateVMSku(context.Background(), oc, resourceSkusClient) utilerror.AssertErrorMessage(t, err, tt.wantErr) diff --git a/pkg/frontend/sku_validation.go b/pkg/frontend/sku_validation.go index 0107dfaa0cb..47bf319e2b0 100644 --- a/pkg/frontend/sku_validation.go +++ b/pkg/frontend/sku_validation.go @@ -42,14 +42,18 @@ func validateVMSku(ctx context.Context, oc *api.OpenShiftCluster, resourceSkusCl // Get a list of available worker SKUs, filtering by location. We initialized a new resourceSkusClient // so that we can determine SKU availability within target cluster subscription instead of within RP subscription. location := oc.Location - filter := fmt.Sprintf("location eq %s", location) - skus, err := resourceSkusClient.List(ctx, filter, false) + + skus := []string{string(oc.Properties.MasterProfile.VMSize)} + workerProfiles, _ := api.GetEnrichedWorkerProfiles(oc.Properties) + for _, workerprofile := range workerProfiles { + skus = append(skus, string(workerprofile.VMSize)) + } + + filteredSkus, err := computeskus.SelectVMSkusInCurrentRegion(ctx, resourceSkusClient, location, skus) if err != nil { return err } - filteredSkus := computeskus.FilterVMSizes(skus, location) - controlPlaneSKU, err := checkSKUAvailability(filteredSkus, location, "properties.masterProfile.VMSize", string(oc.Properties.MasterProfile.VMSize)) if err != nil { return err @@ -67,8 +71,6 @@ func validateVMSku(ctx context.Context, oc *api.OpenShiftCluster, resourceSkusCl } } - workerProfiles, _ := api.GetEnrichedWorkerProfiles(oc.Properties) - // In case there are multiple WorkerProfiles listed in the cluster document (such as post-install), // compare VMSize in each WorkerProfile to the resourceSkusClient call above to ensure that the sku is available in region. // XXX: Will this ever be called post-install? diff --git a/pkg/util/mocks/adminactions/azureactions.go b/pkg/util/mocks/adminactions/azureactions.go index 174227938bc..f97c7d5ca02 100644 --- a/pkg/util/mocks/adminactions/azureactions.go +++ b/pkg/util/mocks/adminactions/azureactions.go @@ -148,6 +148,21 @@ func (mr *MockAzureActionsMockRecorder) ResourcesList(ctx, resources, writer any return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResourcesList", reflect.TypeOf((*MockAzureActions)(nil).ResourcesList), ctx, resources, writer) } +// VMGetSKUs mocks base method. +func (m *MockAzureActions) VMGetSKUs(ctx context.Context, vmSizes []string) (map[string]*armcompute.ResourceSKU, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "VMGetSKUs", ctx, vmSizes) + ret0, _ := ret[0].(map[string]*armcompute.ResourceSKU) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// VMGetSKUs indicates an expected call of VMGetSKUs. +func (mr *MockAzureActionsMockRecorder) VMGetSKUs(ctx, vmSizes any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "VMGetSKUs", reflect.TypeOf((*MockAzureActions)(nil).VMGetSKUs), ctx, vmSizes) +} + // VMRedeployAndWait mocks base method. func (m *MockAzureActions) VMRedeployAndWait(ctx context.Context, vmName string) error { m.ctrl.T.Helper() @@ -191,10 +206,10 @@ func (mr *MockAzureActionsMockRecorder) VMSerialConsole(ctx, log, vmName, target } // VMSizeList mocks base method. -func (m *MockAzureActions) VMSizeList(ctx context.Context) ([]*armcompute.ResourceSKU, error) { +func (m *MockAzureActions) VMSizeList(ctx context.Context) ([]string, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "VMSizeList", ctx) - ret0, _ := ret[0].([]*armcompute.ResourceSKU) + ret0, _ := ret[0].([]string) ret1, _ := ret[1].(error) return ret0, ret1 } diff --git a/pkg/util/mocks/azureclient/azuresdk/armcompute/armcompute.go b/pkg/util/mocks/azureclient/azuresdk/armcompute/armcompute.go index d16b924be52..d4f1272a33d 100644 --- a/pkg/util/mocks/azureclient/azuresdk/armcompute/armcompute.go +++ b/pkg/util/mocks/azureclient/azuresdk/armcompute/armcompute.go @@ -11,6 +11,7 @@ package mock_armcompute import ( context "context" + iter "iter" reflect "reflect" gomock "go.uber.org/mock/gomock" @@ -43,12 +44,11 @@ func (m *MockResourceSKUsClient) EXPECT() *MockResourceSKUsClientMockRecorder { } // List mocks base method. -func (m *MockResourceSKUsClient) List(ctx context.Context, filter string, includeExtendedLocations bool) ([]*armcompute.ResourceSKU, error) { +func (m *MockResourceSKUsClient) List(ctx context.Context, filter string, includeExtendedLocations bool) iter.Seq2[*armcompute.ResourceSKU, error] { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "List", ctx, filter, includeExtendedLocations) - ret0, _ := ret[0].([]*armcompute.ResourceSKU) - ret1, _ := ret[1].(error) - return ret0, ret1 + ret0, _ := ret[0].(iter.Seq2[*armcompute.ResourceSKU, error]) + return ret0 } // List indicates an expected call of List. From 4e65a1aeda8c6664c02ad7e00d2fb42cf8e4aa83 Mon Sep 17 00:00:00 2001 From: Amber Brown Date: Thu, 26 Mar 2026 15:24:46 +1100 Subject: [PATCH 4/5] equalfold and nilchecks --- pkg/util/computeskus/computeskus.go | 21 ++++--- pkg/util/computeskus/computeskus_test.go | 73 +++++++++++++++++++++++- 2 files changed, 83 insertions(+), 11 deletions(-) diff --git a/pkg/util/computeskus/computeskus.go b/pkg/util/computeskus/computeskus.go index 50817338982..5aa318eafad 100644 --- a/pkg/util/computeskus/computeskus.go +++ b/pkg/util/computeskus/computeskus.go @@ -8,6 +8,7 @@ import ( "errors" "fmt" "slices" + "strings" sdkcompute "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v7" @@ -55,9 +56,11 @@ func HasCapability(sku *sdkcompute.ResourceSKU, capabilityName string) bool { // IsRestricted checks whether given resource SKU is restricted in a given location func IsRestricted(sku *sdkcompute.ResourceSKU, location string) bool { for _, restriction := range sku.Restrictions { - for _, restrictedLocation := range restriction.RestrictionInfo.Locations { - if *restrictedLocation == location { - return true + if restriction.RestrictionInfo != nil { + for _, restrictedLocation := range restriction.RestrictionInfo.Locations { + if restrictedLocation != nil && strings.EqualFold(*restrictedLocation, location) { + return true + } } } } @@ -92,12 +95,12 @@ func SelectVMSkusInCurrentRegion(ctx context.Context, resourceSkusClient armcomp continue } - if slices.Contains(skuNames, *sku.Name) { - // Make sure it's actually in our location - if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return *s == location }) { - continue - } + // Make sure it's actually in our location + if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return s != nil && strings.EqualFold(*s, location) }) { + continue + } + if slices.Contains(skuNames, *sku.Name) { vmskus[*sku.Name] = sku } @@ -126,7 +129,7 @@ func ListUnrestrictedVMSkusInCurrentRegion(ctx context.Context, resourceSkusClie } // Make sure it's actually in our location - if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return *s == location }) { + if !slices.ContainsFunc(sku.Locations, func(s *string) bool { return s != nil && strings.EqualFold(*s, location) }) { continue } diff --git a/pkg/util/computeskus/computeskus_test.go b/pkg/util/computeskus/computeskus_test.go index bc70228965e..5db3cc316d4 100644 --- a/pkg/util/computeskus/computeskus_test.go +++ b/pkg/util/computeskus/computeskus_test.go @@ -254,6 +254,17 @@ func TestSelectVMSkusInCurrentRegion(t *testing.T) { }, }, }: nil, + // Nil region + { + Name: pointerutils.ToPtr("smallmachine_v6"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: []*string{nil}, + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, // Machine that has no locations/locationinfo { Name: pointerutils.ToPtr("smallmachine_v3"), @@ -333,7 +344,6 @@ func TestSelectVMSkusInCurrentRegion(t *testing.T) { Location: pointerutils.ToPtr("northus2"), }, }, - Restrictions: []*armcompute.ResourceSKURestrictions{}, }: nil, // Machine that has no locations/locationinfo { @@ -460,6 +470,17 @@ func TestListUnrestrictedSKUNames(t *testing.T) { }, }, }: nil, + // Capitalisation of region + { + Name: pointerutils.ToPtr("smallmachine_v10"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("Northus2"), + }, + }, + }: nil, // Duplicated struct, in case we get two { Name: pointerutils.ToPtr("smallmachine_v4"), @@ -483,6 +504,17 @@ func TestListUnrestrictedSKUNames(t *testing.T) { }, }, }: nil, + // Nil region + { + Name: pointerutils.ToPtr("smallmachine_v11"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: []*string{nil}, + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }: nil, // Restricted in this region { Name: pointerutils.ToPtr("smallmachine_v2"), @@ -494,6 +526,11 @@ func TestListUnrestrictedSKUNames(t *testing.T) { }, }, Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"somewhereelse"}), + }, + }, { RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ Locations: pointerutils.ToSlicePtr([]string{"northus2"}), @@ -501,6 +538,38 @@ func TestListUnrestrictedSKUNames(t *testing.T) { }, }, }: nil, + // Restricted in this region, equal fold + { + Name: pointerutils.ToPtr("smallmachine_v20"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + }, + }, + }, + }: nil, + // Nil restriction info + { + Name: pointerutils.ToPtr("smallmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + {}, + }, + }: nil, // Machine that has no locations/locationinfo { Name: pointerutils.ToPtr("smallmachine_v3"), @@ -522,7 +591,7 @@ func TestListUnrestrictedSKUNames(t *testing.T) { }), ) }, - desired: []string{"bigmachine_v1", "smallmachine_v4"}, + desired: []string{"bigmachine_v1", "smallmachine_v1", "smallmachine_v10", "smallmachine_v4"}, }, { name: "duplicate VM structs don't lead to duplicated names", From 67d6fb60e8cb617f6e16de083d7c8ec7ae670f8e Mon Sep 17 00:00:00 2001 From: Amber Brown Date: Fri, 27 Mar 2026 12:20:40 +1100 Subject: [PATCH 5/5] change tests so that copilot can understand them, some more nil handling --- pkg/util/computeskus/computeskus.go | 10 +- pkg/util/computeskus/computeskus_test.go | 736 ++++++++++++----------- 2 files changed, 404 insertions(+), 342 deletions(-) diff --git a/pkg/util/computeskus/computeskus.go b/pkg/util/computeskus/computeskus.go index 5aa318eafad..4a3f31398dd 100644 --- a/pkg/util/computeskus/computeskus.go +++ b/pkg/util/computeskus/computeskus.go @@ -56,7 +56,7 @@ func HasCapability(sku *sdkcompute.ResourceSKU, capabilityName string) bool { // IsRestricted checks whether given resource SKU is restricted in a given location func IsRestricted(sku *sdkcompute.ResourceSKU, location string) bool { for _, restriction := range sku.Restrictions { - if restriction.RestrictionInfo != nil { + if restriction != nil && restriction.RestrictionInfo != nil { for _, restrictedLocation := range restriction.RestrictionInfo.Locations { if restrictedLocation != nil && strings.EqualFold(*restrictedLocation, location) { return true @@ -90,6 +90,10 @@ func SelectVMSkusInCurrentRegion(ctx context.Context, resourceSkusClient armcomp return nil, fmt.Errorf("%w: %w", ErrListVMResourceSKUs, err) } + if sku == nil { + continue + } + // We only care about VMs and ones with locations/locationinfo if *sku.ResourceType != "virtualMachines" || len(sku.Locations) == 0 || len(sku.LocationInfo) == 0 { continue @@ -123,6 +127,10 @@ func ListUnrestrictedVMSkusInCurrentRegion(ctx context.Context, resourceSkusClie return nil, fmt.Errorf("%w: %w", ErrListVMResourceSKUs, err) } + if sku == nil { + continue + } + // We only care about VMs and ones with locations/locationinfo if *sku.ResourceType != "virtualMachines" || len(sku.Locations) == 0 || len(sku.LocationInfo) == 0 { continue diff --git a/pkg/util/computeskus/computeskus_test.go b/pkg/util/computeskus/computeskus_test.go index 5db3cc316d4..62e4ecd5d9e 100644 --- a/pkg/util/computeskus/computeskus_test.go +++ b/pkg/util/computeskus/computeskus_test.go @@ -5,7 +5,6 @@ package computeskus import ( "errors" - "maps" "reflect" "testing" @@ -221,70 +220,144 @@ func TestSelectVMSkusInCurrentRegion(t *testing.T) { name: "happypath", vmSkus: []string{"bigmachine_v1", "smallmachine_v4", "smallmachine_v5"}, mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { - mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - // Not actually in our region - { - Name: pointerutils.ToPtr("smallmachine_v5"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus1"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus1"), + // Capitalisation of region + { + Name: pointerutils.ToPtr("smallmachine_v10"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("Northus2"), + }, }, }, - }: nil, - // Nil region - { - Name: pointerutils.ToPtr("smallmachine_v6"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: []*string{nil}, - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Not actually in our region + { + Name: pointerutils.ToPtr("smallmachine_v9"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus1"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus1"), + }, }, }, - }: nil, - // Machine that has no locations/locationinfo - { - Name: pointerutils.ToPtr("smallmachine_v3"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: nil, - LocationInfo: nil, - }: nil, - // Actually an availabilitySet - { - Name: pointerutils.ToPtr("Classic"), - ResourceType: pointerutils.ToPtr("availabilitySets"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Nil region + { + Name: pointerutils.ToPtr("smallmachine_v11"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: []*string{nil}, + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"somewhereelse"}), + }, + }, + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, + }, + }, + }, + // Restricted in this region, equal fold + { + Name: pointerutils.ToPtr("smallmachine_v20"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + }, + }, + }, + }, + // An actual nil + nil, + // Nil restriction info + { + Name: pointerutils.ToPtr("smallmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + Restrictions: []*armcompute.ResourceSKURestrictions{ + {}, + nil, + }, + }, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - }), - ) + } { + if !yield(s, nil) { + return + } + } + }) }, desired: map[string]*armcompute.ResourceSKU{ "bigmachine_v1": { @@ -310,61 +383,38 @@ func TestSelectVMSkusInCurrentRegion(t *testing.T) { }, }, { - name: "duplicate skus in vmskus", + name: "duplicate skus in vmskus input list does not error", vmSkus: []string{"bigmachine_v1", "bigmachine_v1", "smallmachine_v4"}, mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { - mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - // Machine that has no locations/locationinfo - { - Name: pointerutils.ToPtr("smallmachine_v3"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: nil, - LocationInfo: nil, - }: nil, - // Actually an availabilitySet - { - Name: pointerutils.ToPtr("Classic"), - ResourceType: pointerutils.ToPtr("availabilitySets"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - }), - ) + } { + if !yield(s, nil) { + return + } + } + }) }, desired: map[string]*armcompute.ResourceSKU{ "bigmachine_v1": { @@ -393,31 +443,36 @@ func TestSelectVMSkusInCurrentRegion(t *testing.T) { name: "error in pagination", vmSkus: []string{"bigmachine_v5"}, mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { - mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - nil: errors.New("this is an error"), - }), - ) + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, + }, + }, + } { + if !yield(s, nil) { + return + } + } + yield(nil, errors.New("this is an error")) + }) }, wantError: "this is an error", }, @@ -448,148 +503,156 @@ func TestListUnrestrictedSKUNames(t *testing.T) { { name: "happypath", mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { - mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - // Capitalisation of region - { - Name: pointerutils.ToPtr("smallmachine_v10"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("Northus2"), + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - // Duplicated struct, in case we get two - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Capitalisation of region + { + Name: pointerutils.ToPtr("smallmachine_v10"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("Northus2"), + }, }, }, - Restrictions: []*armcompute.ResourceSKURestrictions{}, - }: nil, - // Not actually in our region - { - Name: pointerutils.ToPtr("smallmachine_v9"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus1"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus1"), + // Duplicated struct, in case we get two + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, + Restrictions: []*armcompute.ResourceSKURestrictions{}, }, - }: nil, - // Nil region - { - Name: pointerutils.ToPtr("smallmachine_v11"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: []*string{nil}, - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Not actually in our region + { + Name: pointerutils.ToPtr("smallmachine_v9"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus1"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus1"), + }, }, }, - }: nil, - // Restricted in this region - { - Name: pointerutils.ToPtr("smallmachine_v2"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Nil region + { + Name: pointerutils.ToPtr("smallmachine_v11"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: []*string{nil}, + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - Restrictions: []*armcompute.ResourceSKURestrictions{ - { - RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ - Locations: pointerutils.ToSlicePtr([]string{"somewhereelse"}), + // An actual nil + nil, + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), }, }, - { - RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"somewhereelse"}), + }, + }, + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, }, }, }, - }: nil, - // Restricted in this region, equal fold - { - Name: pointerutils.ToPtr("smallmachine_v20"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Restricted in this region, equal fold + { + Name: pointerutils.ToPtr("smallmachine_v20"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, - }, - Restrictions: []*armcompute.ResourceSKURestrictions{ - { - RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ - Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"Northus2"}), + }, }, }, }, - }: nil, - // Nil restriction info - { - Name: pointerutils.ToPtr("smallmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Nil restriction info + { + Name: pointerutils.ToPtr("smallmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, - }, - Restrictions: []*armcompute.ResourceSKURestrictions{ - {}, - }, - }: nil, - // Machine that has no locations/locationinfo - { - Name: pointerutils.ToPtr("smallmachine_v3"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: nil, - LocationInfo: nil, - }: nil, - // Actually an availabilitySet - { - Name: pointerutils.ToPtr("Classic"), - ResourceType: pointerutils.ToPtr("availabilitySets"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + Restrictions: []*armcompute.ResourceSKURestrictions{ + {}, + nil, + }, + }, + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - }), - ) + } { + if !yield(s, nil) { + return + } + } + }) }, desired: []string{"bigmachine_v1", "smallmachine_v1", "smallmachine_v10", "smallmachine_v4"}, }, @@ -597,125 +660,116 @@ func TestListUnrestrictedSKUNames(t *testing.T) { name: "duplicate VM structs don't lead to duplicated names", mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, - {}, }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - // Restricted in this region - { - Name: pointerutils.ToPtr("smallmachine_v2"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Restricted in this region + { + Name: pointerutils.ToPtr("smallmachine_v2"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, - }, - Restrictions: []*armcompute.ResourceSKURestrictions{ - { - RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + Restrictions: []*armcompute.ResourceSKURestrictions{ + { + RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + }, }, }, }, - }: nil, - // Machine that has no locations/locationinfo - { - Name: pointerutils.ToPtr("smallmachine_v3"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: nil, - LocationInfo: nil, - }: nil, - // Actually an availabilitySet - { - Name: pointerutils.ToPtr("Classic"), - ResourceType: pointerutils.ToPtr("availabilitySets"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + // Machine that has no locations/locationinfo + { + Name: pointerutils.ToPtr("smallmachine_v3"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: nil, + LocationInfo: nil, + }, + // Actually an availabilitySet + { + Name: pointerutils.ToPtr("Classic"), + ResourceType: pointerutils.ToPtr("availabilitySets"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - }: nil, - }), - ) + } { + if !yield(s, nil) { + return + } + } + }) }, desired: []string{"bigmachine_v1", "smallmachine_v4"}, }, { name: "error in pagination", mocks: func(mrsc *mock_armcompute.MockResourceSKUsClient) { - mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false).Return( - maps.All(map[*armcompute.ResourceSKU]error{ - { - Name: pointerutils.ToPtr("bigmachine_v1"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - { - Name: pointerutils.ToPtr("smallmachine_v4"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), - }, - }, - }: nil, - // Restricted in this region - { - Name: pointerutils.ToPtr("smallmachine_v2"), - ResourceType: pointerutils.ToPtr("virtualMachines"), - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), - LocationInfo: []*armcompute.ResourceSKULocationInfo{ - { - Location: pointerutils.ToPtr("northus2"), + mrsc.EXPECT().List(gomock.Any(), "location eq northus2", false). + Return(func(yield func(*armcompute.ResourceSKU, error) bool) { + for _, s := range []*armcompute.ResourceSKU{ + { + Name: pointerutils.ToPtr("bigmachine_v1"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), + }, }, }, - Restrictions: []*armcompute.ResourceSKURestrictions{ - { - RestrictionInfo: &armcompute.ResourceSKURestrictionInfo{ - Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + { + Name: pointerutils.ToPtr("smallmachine_v4"), + ResourceType: pointerutils.ToPtr("virtualMachines"), + Locations: pointerutils.ToSlicePtr([]string{"northus2"}), + LocationInfo: []*armcompute.ResourceSKULocationInfo{ + { + Location: pointerutils.ToPtr("northus2"), }, }, }, - }: nil, - nil: errors.New("this is an error"), - }), - ) + } { + if !yield(s, nil) { + return + } + } + yield(nil, errors.New("this is an error")) + }) }, wantError: "this is an error", },