From 93679382f13eda30c24c63da97ba859f3751f8aa Mon Sep 17 00:00:00 2001
From: Yaroslava Serdiuk <yaroslava@google.com>
Date: Tue, 4 Jun 2024 09:34:44 +0000
Subject: [PATCH] Public ProvisioningClass interface

---
 cluster-autoscaler/main.go                    | 11 +++++----
 .../orchestrator/orchestrator.go              | 24 ++++++-------------
 .../orchestrator/orchestrator_test.go         |  2 +-
 3 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/cluster-autoscaler/main.go b/cluster-autoscaler/main.go
index c5ff21aaa4e7..c59b1a50b724 100644
--- a/cluster-autoscaler/main.go
+++ b/cluster-autoscaler/main.go
@@ -32,6 +32,7 @@ import (
 	"k8s.io/autoscaler/cluster-autoscaler/core/scaleup/orchestrator"
 	"k8s.io/autoscaler/cluster-autoscaler/debuggingsnapshot"
 	"k8s.io/autoscaler/cluster-autoscaler/loop"
+	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/besteffortatomic"
 	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/checkcapacity"
 	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/provreqclient"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/predicatechecker"
@@ -496,17 +497,17 @@ func buildAutoscaler(debuggingSnapshotter debuggingsnapshot.DebuggingSnapshotter
 		podListProcessor.AddProcessor(provreq.NewProvisioningRequestPodsFilter(provreq.NewDefautlEventManager()))
 
 		restConfig := kube_util.GetKubeConfig(autoscalingOptions.KubeClientOpts)
-		provreqOrchestrator, err := provreqorchestrator.New(restConfig)
+		client, err := provreqclient.NewProvisioningRequestClient(restConfig)
 		if err != nil {
 			return nil, err
 		}
+		provreqOrchestrator := provreqorchestrator.New(client, []provreqorchestrator.ProvisioningClass{
+			checkcapacity.New(client),
+			besteffortatomic.New(client),
+		})
 		scaleUpOrchestrator := provreqorchestrator.NewWrapperOrchestrator(provreqOrchestrator)
 
 		opts.ScaleUpOrchestrator = scaleUpOrchestrator
-		client, err := provreqclient.NewProvisioningRequestClient(restConfig)
-		if err != nil {
-			return nil, err
-		}
 		provreqProcesor := provreq.NewCombinedProvReqProcessor(client, []provreq.ProvisioningRequestProcessor{checkcapacity.NewCheckCapacityProcessor(client)})
 		if err != nil {
 			return nil, err
diff --git a/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator.go b/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator.go
index ab86f7193cf7..a2f7b9116fc3 100644
--- a/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator.go
+++ b/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator.go
@@ -27,22 +27,20 @@ import (
 	"k8s.io/autoscaler/cluster-autoscaler/context"
 	"k8s.io/autoscaler/cluster-autoscaler/estimator"
 	"k8s.io/autoscaler/cluster-autoscaler/processors/status"
-	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/besteffortatomic"
-	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/checkcapacity"
 	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/conditions"
 	provreq_pods "k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/pods"
 	"k8s.io/autoscaler/cluster-autoscaler/provisioningrequest/provreqclient"
 	"k8s.io/autoscaler/cluster-autoscaler/simulator/scheduling"
 	ca_errors "k8s.io/autoscaler/cluster-autoscaler/utils/errors"
 	"k8s.io/autoscaler/cluster-autoscaler/utils/taints"
-	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
 
 	ca_processors "k8s.io/autoscaler/cluster-autoscaler/processors"
 	schedulerframework "k8s.io/kubernetes/pkg/scheduler/framework"
 )
 
-type provisioningClass interface {
+// ProvisioningClass is an interface for ProvisioningRequests.
+type ProvisioningClass interface {
 	Provision([]*apiv1.Pod, []*apiv1.Node, []*appsv1.DaemonSet,
 		map[string]*schedulerframework.NodeInfo) (*status.ScaleUpStatus, ca_errors.AutoscalerError)
 	Initialize(*context.AutoscalingContext, *ca_processors.AutoscalingProcessors, *clusterstate.ClusterStateRegistry,
@@ -55,23 +53,15 @@ type provReqOrchestrator struct {
 	context             *context.AutoscalingContext
 	client              *provreqclient.ProvisioningRequestClient
 	injector            *scheduling.HintingSimulator
-	provisioningClasses []provisioningClass
+	provisioningClasses []ProvisioningClass
 }
 
 // New return new orchestrator.
-func New(kubeConfig *rest.Config) (*provReqOrchestrator, error) {
-	client, err := provreqclient.NewProvisioningRequestClient(kubeConfig)
-	if err != nil {
-		return nil, err
-	}
-
+func New(client *provreqclient.ProvisioningRequestClient, classes []ProvisioningClass) *provReqOrchestrator {
 	return &provReqOrchestrator{
-		client: client,
-		provisioningClasses: []provisioningClass{
-			checkcapacity.New(client),
-			besteffortatomic.New(client),
-		},
-	}, nil
+		client:              client,
+		provisioningClasses: classes,
+	}
 }
 
 // Initialize initialize orchestrator.
diff --git a/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator_test.go b/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator_test.go
index fa2fd55576c3..6e0ead7d6283 100644
--- a/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator_test.go
+++ b/cluster-autoscaler/provisioningrequest/orchestrator/orchestrator_test.go
@@ -334,7 +334,7 @@ func setupTest(t *testing.T, nodes []*apiv1.Node, prs []*provreqwrapper.Provisio
 
 	orchestrator := &provReqOrchestrator{
 		client:              client,
-		provisioningClasses: []provisioningClass{checkcapacity.New(client), besteffortatomic.New(client)},
+		provisioningClasses: []ProvisioningClass{checkcapacity.New(client), besteffortatomic.New(client)},
 	}
 	orchestrator.Initialize(&autoscalingContext, processors, clusterState, estimatorBuilder, taints.TaintConfig{})
 	return orchestrator, nodeInfos