address comments and fix test

Signed-off-by: Britania Rodriguez Reyes <britaniar@microsoft.com>

Author: britaniar

charts/hub-agent/templates/deployment.yaml

@@ -54,7 +54,7 @@ spec:
             - --cluster-unhealthy-threshold={{ .Values.clusterUnhealthyThreshold }}
             - --resource-snapshot-creation-minimum-interval={{ .Values.resourceSnapshotCreationMinimumInterval }}
             - --resource-changes-collection-duration={{ .Values.resourceChangesCollectionDuration }}
-            - --compute-endpoint={{ .Values.computeEndpoint }}
+            - --compute-service-endpoint={{ .Values.computeServiceEndpoint }}
           ports:
             - name: metrics
               containerPort: 8080

charts/hub-agent/values.yaml

@@ -54,7 +54,7 @@ enableEvictionAPIs: true
 enablePprof: true
 pprofPort: 6065
 
-computeEndpoint: "http://localhost:9090"
+computeServiceEndpoint: "http://localhost:8421"
 
 hubAPIQPS: 250
 hubAPIBurst: 1000

cmd/hubagent/options/options.go

@@ -104,8 +104,8 @@ type Options struct {
 	EnablePprof bool
 	// PprofPort is the port for pprof profiling.
 	PprofPort int
-	// ComputeEndpoint is the endpoint for the compute service.
-	ComputeEndpoint string
+	// ComputeServiceEndpoint is the endpoint for the compute service.
+	ComputeServiceEndpoint string
 	// DenyModifyMemberClusterLabels indicates if the member cluster labels cannot be modified by groups (excluding system:masters)
 	DenyModifyMemberClusterLabels bool
 	// ResourceSnapshotCreationMinimumInterval is the minimum interval at which resource snapshots could be created.
@@ -133,7 +133,7 @@ func NewOptions() *Options {
 		EnableResourcePlacement:                 true,
 		EnablePprof:                             false,
 		PprofPort:                               6065,
-		ComputeEndpoint:                         "http://localhost:9090",
+		ComputeServiceEndpoint:                  "http://localhost:8421",
 		ResourceSnapshotCreationMinimumInterval: 30 * time.Second,
 		ResourceChangesCollectionDuration:       15 * time.Second,
 	}
@@ -182,7 +182,7 @@ func (o *Options) AddFlags(flags *flag.FlagSet) {
 	flags.BoolVar(&o.EnableResourcePlacement, "enable-resource-placement", true, "If set, the agents will watch for the ResourcePlacement APIs.")
 	flags.BoolVar(&o.EnablePprof, "enable-pprof", false, "If set, the pprof profiling is enabled.")
 	flags.IntVar(&o.PprofPort, "pprof-port", 6065, "The port for pprof profiling.")
-	flags.StringVar(&o.ComputeEndpoint, "compute-endpoint", "http://localhost:9090", "The endpoint for the compute service.")
+	flags.StringVar(&o.ComputeServiceEndpoint, "compute-service-endpoint", "http://localhost:8421", "The endpoint for the compute service.")
 	flags.BoolVar(&o.DenyModifyMemberClusterLabels, "deny-modify-member-cluster-labels", false, "If set, users not in the system:masters cannot modify member cluster labels.")
 	flags.DurationVar(&o.ResourceSnapshotCreationMinimumInterval, "resource-snapshot-creation-minimum-interval", 30*time.Second, "The minimum interval at which resource snapshots could be created.")
 	flags.DurationVar(&o.ResourceChangesCollectionDuration, "resource-changes-collection-duration", 15*time.Second,

cmd/hubagent/workload/setup.go

@@ -22,6 +22,8 @@ import (
 	"strings"
 	"sync"
 
+	"go.goms.io/fleet/pkg/scheduler/framework/plugins/clusteraffinity/azure"
+
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/client-go/discovery"
 	"k8s.io/client-go/dynamic"
@@ -385,7 +387,7 @@ func SetupControllers(ctx context.Context, wg *sync.WaitGroup, mgr ctrl.Manager,
 		// Set up the scheduler
 		klog.Info("Setting up scheduler")
 		defaultProfile := profile.NewDefaultProfileWithOptions(profile.ProfileOptions{
-			Endpoint: opts.ComputeEndpoint,
+			ComputeService: azure.NewAzureCapacityService(opts.ComputeServiceEndpoint),
 		})
 		defaultFramework := framework.NewFramework(defaultProfile, mgr)
 		defaultSchedulingQueue := queue.NewSimplePlacementSchedulingQueue(

pkg/scheduler/framework/plugins/clusteraffinity/azure/capcacity.go

@@ -1,7 +1,7 @@
 package azure
 
 import (
-	"context"
+	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -15,29 +15,34 @@ import (
 	placementv1beta1 "go.goms.io/fleet/apis/placement/v1beta1"
 )
 
-// DefaultAzureCapacityService is the default implementation of AzureCapacityService.
-type DefaultAzureCapacityService struct {
-	endpoint string
-	client   *http.Client
+// defaultAzureCapacityService is the default implementation of AzureCapacityService.
+type defaultAzureCapacityService struct {
+	// ServerAddress is the Azure capacity service endpoint.
+	ServerAddress string
+
+	// Client is the HTTP client used to make requests to the Azure capacity service.
+	Client *http.Client
 }
 
-// Compile-time check to ensure DefaultAzureCapacityService implements AzureCapacityService
-var _ AzureCapacityService = &DefaultAzureCapacityService{}
+// Compile-time check to ensure defaultAzureCapacityService implements AzureCapacityService
+var _ AzureCapacityService = &defaultAzureCapacityService{}
 
 // NewAzureCapacityService creates a new default Azure capacity service with the given endpoint.
-func NewAzureCapacityService(endpoint string) *DefaultAzureCapacityService {
-	return &DefaultAzureCapacityService{
-		endpoint: endpoint,
-		client:   &http.Client{},
+func NewAzureCapacityService(serverAddress string, client *http.Client) AzureCapacityService {
+	if client == nil {
+		client = http.DefaultClient
+	}
+	return &defaultAzureCapacityService{
+		ServerAddress: serverAddress,
+		Client:        client,
 	}
 }
 
 // ValidateCapacityRequirement validates a capacity requirement against Azure APIs.
-func (s *DefaultAzureCapacityService) ValidateCapacityRequirement(
+func (s *defaultAzureCapacityService) ValidateCapacityRequirement(
 	cluster *clusterv1beta1.MemberCluster,
 	req placementv1beta1.PropertySelectorRequirement,
 ) (bool, error) {
-	ctx := context.Background()
 	subID, location := extractAzureInfoFromLabels(cluster)
 	if subID == "" || location == "" {
 		return false, fmt.Errorf("cluster %s does not have required Azure labels", cluster.Name)
@@ -48,7 +53,7 @@ func (s *DefaultAzureCapacityService) ValidateCapacityRequirement(
 	}
 
 	// TODO: Replace with actual Azure capacity client call
-	endpoint := s.endpoint
+	endpoint := s.ServerAddress
 	if !strings.HasPrefix(endpoint, "http://") && !strings.HasPrefix(endpoint, "https://") {
 		endpoint = "http://" + endpoint
 	}
@@ -58,38 +63,41 @@ func (s *DefaultAzureCapacityService) ValidateCapacityRequirement(
 		"subscription_id": subID,
 		"location":        location,
 		"regular_priority_profile": map[string]interface{}{
-			"capacity_unit_type": 1, // CAPACITY_UNIT_TYPE_VM_INSTANCE_COUNT
-			"target_capacity":    uint32(capacity.Value()),
+			"capacity_unit_type": 1,                       // CAPACITY_UNIT_TYPE_VM_INSTANCE_COUNT
+			"target_capacity":    int32(capacity.Value()), // Safe conversion as VM counts won't exceed int32
 		},
 		"recommendation_properties": map[string]interface{}{
 			"restrictions_filter": 4, // RESTRICTIONS_FILTER_QUOTA_AND_OFFER_RESTRICTIONS
 		},
 		"resource_properties": map[string]interface{}{},
 	}
-	body, err := json.Marshal(payload)
+	// Convert the payload to JSON
+	jsonData, err := json.Marshal(payload)
 	if err != nil {
 		return false, fmt.Errorf("failed to marshal request payload: %w", err)
 	}
 
-	httpReq, err := http.NewRequestWithContext(ctx, "POST", url, strings.NewReader(string(body)))
+	// Create a new POST request
+	httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(jsonData))
 	if err != nil {
 		return false, fmt.Errorf("failed to create HTTP request: %w", err)
 	}
+
 	httpReq.Header.Set("Content-Type", "application/json")
 	httpReq.Header.Set("Accept", "application/json")
 	// Add service mesh headers
 	httpReq.Header.Set("x-service-name", "fleet")
 	httpReq.Header.Set("x-service-version", "v1")
 
-	resp, err := s.client.Do(httpReq)
+	resp, err := s.Client.Do(httpReq)
 	if err != nil {
 		return false, fmt.Errorf("failed to make HTTP request to Azure service: %w", err)
 	}
 	defer resp.Body.Close()
 
 	bodyBytes, _ := io.ReadAll(resp.Body)
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		return false, fmt.Errorf("Azure service returned status %d: %s", resp.StatusCode, string(bodyBytes))
+		return false, fmt.Errorf("azure service returned status %d: %s", resp.StatusCode, string(bodyBytes))
 	}
 
 	var respObj struct {
@@ -143,14 +151,14 @@ func extractCapacityRequirements(req placementv1beta1.PropertySelectorRequiremen
 	}
 
 	// Remove prefix and suffix to get the SKU
-	sku := strings.TrimSuffix(strings.TrimPrefix(req.Name, SkuCapacityPropertyPrefix+"."), "/capacity")
+	sku := strings.TrimSuffix(strings.TrimPrefix(req.Name, SkuCapacityPropertyPrefix+"/"), "/capacity")
 	if sku == "" {
 		return nil, "", fmt.Errorf("cannot extract SKU from property name: %q", req.Name)
 	}
 
 	// Validate that we have exactly one value
 	if len(req.Values) != 1 {
-		return nil, "", fmt.Errorf("Azure SKU capacity property must have exactly one value, got %d", len(req.Values))
+		return nil, "", fmt.Errorf("azure SKU capacity property must have exactly one value, got %d", len(req.Values))
 	}
 
 	// Parse the capacity value

pkg/scheduler/framework/plugins/clusteraffinity/azure/capcacity_test.go

@@ -1,17 +1,17 @@
 package azure
 
 import (
-	"encoding/json"
 	"net/http"
 	"net/http/httptest"
 	"strings"
 	"testing"
-
-	"k8s.io/apimachinery/pkg/api/resource"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"time"
 
 	clusterv1beta1 "go.goms.io/fleet/apis/cluster/v1beta1"
 	placementv1beta1 "go.goms.io/fleet/apis/placement/v1beta1"
+
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 func TestExtractAzureInfoFromLabels(t *testing.T) {
@@ -239,31 +239,85 @@ func TestExtractCapacityRequirements(t *testing.T) {
 	}
 }
 
-func TestValidateCapacityRequirement(t *testing.T) {
-	// Mock response JSON for recommendedVmSizes
-	mockResponse := map[string]interface{}{
-		"recommendedVmSizes": map[string]interface{}{
-			"regularVmSizes": []map[string]interface{}{
-				{"name": "Standard_D2s_v3"},
-				{"name": "Standard_A1_v2"},
-			},
+func TestNewAzureCapacityService(t *testing.T) {
+	tests := []struct {
+		name          string
+		serverAddress string
+		client        *http.Client
+	}{
+		{
+			name:          "with valid client",
+			serverAddress: "test.example.com",
+			client:        &http.Client{},
+		},
+		{
+			name:          "with nil client",
+			serverAddress: "test.example.com",
+			client:        nil,
 		},
 	}
-	respBytes, _ := json.Marshal(mockResponse)
 
-	// Start a test HTTP server
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Set("Content-Type", "application/json")
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			svc := NewAzureCapacityService(tt.serverAddress, tt.client).(*defaultAzureCapacityService)
+			if svc == nil {
+				t.Fatal("NewAzureCapacityService returned nil")
+			}
+			if svc.Client == nil {
+				t.Error("client is nil")
+			}
+			if svc.ServerAddress != tt.serverAddress {
+				t.Errorf("serverAddress = %v, want %v", svc.ServerAddress, tt.serverAddress)
+			}
+		})
+	}
+}
+
+// Mock the expected response from the Azure API.
+var mockAzureResponse = `{
+    "recommendedVmSizes": {
+        "regularVmSizes": [
+            {
+                "family": "Dsv3",
+                "name": "Standard_D2s_v3",
+                "size": "D2"
+            },
+            {
+                "family": "Standard",
+                "name": "Standard_B1s",
+                "size": "Standard_B1s"
+            }
+        ]
+    }
+}`
+
+func TestValidateCapacityRequirement(t *testing.T) {
+	// Create mock server
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Verify request method
+		if r.Method != http.MethodPost {
+			t.Errorf("got %s, want POST request", r.Method)
+		}
+
+		// Verify headers
+		if r.Header.Get("Content-Type") != "application/json" {
+			t.Errorf("got %s, want Content-Type: application/json", r.Header.Get("Content-Type"))
+		}
+		if r.Header.Get("Accept") != "application/json" {
+			t.Errorf("got %s, want Accept: application/json", r.Header.Get("Accept"))
+		}
+
+		// Write mock response
 		w.WriteHeader(http.StatusOK)
-		w.Write(respBytes)
+		if _, err := w.Write([]byte(mockAzureResponse)); err != nil {
+			t.Fatalf("failed to write response: %v", err)
+		}
 	}))
-	defer ts.Close()
 
-	// Create the service pointing to the test server
-	svc := &DefaultAzureCapacityService{
-		endpoint: ts.URL,
-		client:   ts.Client(),
-	}
+	client := &http.Client{Timeout: 60 * time.Second}
+
+	// Step 2: Create the AzureCapacityService with the mock client and server URL
+	svc := NewAzureCapacityService(server.URL, client)
 
 	// Prepare test data
 	cluster := &clusterv1beta1.MemberCluster{
@@ -279,6 +333,7 @@ func TestValidateCapacityRequirement(t *testing.T) {
 		name           string
 		cluster        *clusterv1beta1.MemberCluster
 		req            placementv1beta1.PropertySelectorRequirement
+		mockStatusCode int
 		wantAvailable  bool
 		expectError    bool
 		errorSubstring string
@@ -287,18 +342,21 @@ func TestValidateCapacityRequirement(t *testing.T) {
 			name:    "valid request",
 			cluster: cluster,
 			req: placementv1beta1.PropertySelectorRequirement{
-				Name:   "kubernetes.azure.com/vm-size/Standard_D2s_v3/capacity",
-				Values: []string{"10"},
+				Name:     "kubernetes.azure.com/vm-size/Standard_D2s_v3/capacity",
+				Operator: placementv1beta1.PropertySelectorGreaterThan,
+				Values:   []string{"10"},
 			},
-			wantAvailable: true,
-			expectError:   false,
+			mockStatusCode: http.StatusOK,
+			wantAvailable:  true,
+			expectError:    false,
 		},
 		{
 			name:    "invalid capacity request",
 			cluster: cluster,
 			req: placementv1beta1.PropertySelectorRequirement{
-				Name:   "kubernetes.azure.com/vm-size/Standard_D3s_v3/capacity",
-				Values: []string{"2"},
+				Name:     "kubernetes.azure.com/vm-size/Standard_D3s_v3/capacity",
+				Operator: placementv1beta1.PropertySelectorGreaterThan,
+				Values:   []string{"2"},
 			},
 			wantAvailable: false,
 			expectError:   false,
@@ -313,8 +371,9 @@ func TestValidateCapacityRequirement(t *testing.T) {
 				},
 			},
 			req: placementv1beta1.PropertySelectorRequirement{
-				Name:   "kubernetes.azure.com/vm-size/Standard_D2s_v3/capacity",
-				Values: []string{"10"},
+				Name:     "kubernetes.azure.com/vm-size/Standard_D2s_v3/capacity",
+				Operator: placementv1beta1.PropertySelectorGreaterThan,
+				Values:   []string{"10"},
 			},
 			expectError:    true,
 			errorSubstring: "does not have required Azure labels",
@@ -323,14 +382,14 @@ func TestValidateCapacityRequirement(t *testing.T) {
 			name:    "invalid capacity property",
 			cluster: cluster,
 			req: placementv1beta1.PropertySelectorRequirement{
-				Name:   "kubernetes.azure.com/vm-size/Standard_D2s_v3/usage",
-				Values: []string{"10"},
+				Name:     "kubernetes.azure.com/vm-size/Standard_D2s_v3/usage",
+				Operator: placementv1beta1.PropertySelectorGreaterThan,
+				Values:   []string{"10"},
 			},
 			expectError:    true,
 			errorSubstring: "invalid Azure SKU capacity property format",
 		},
 	}
-
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			result, err := svc.ValidateCapacityRequirement(tt.cluster, tt.req)