APIGOV-28984 Updates

Axway · Nov 5, 2024 · 0b424b3 · 0b424b3
1 parent 6f49a1d
commit 0b424b3
Show file tree

Hide file tree

Showing 8 changed files with 220 additions and 159 deletions.
diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
@@ -25,7 +25,6 @@ import (
 	"github.com/Axway/agent-sdk/pkg/cache"
 	"github.com/Axway/agent-sdk/pkg/config"
 	"github.com/Axway/agent-sdk/pkg/customunit"
-	customunithandler "github.com/Axway/agent-sdk/pkg/customunit/handler"
 	"github.com/Axway/agent-sdk/pkg/util"
 	hc "github.com/Axway/agent-sdk/pkg/util/healthcheck"
 	"github.com/Axway/agent-sdk/pkg/util/log"
@@ -63,17 +62,17 @@ type agentData struct {
 	agentFeaturesCfg     config.AgentFeaturesConfig
 	tokenRequester       auth.PlatformTokenGetter
 
-	teamMap                    cache.Cache
-	cacheManager               agentcache.Manager
-	apiValidator               APIValidator
-	apiValidatorLock           sync.Mutex
-	apiValidatorJobID          string
-	configChangeHandler        ConfigChangeHandler
-	agentResourceChangeHandler ConfigChangeHandler
-	customUnitQuotaHandler     customunithandler.CustomUnitQuotaHandler
-	agentShutdownHandler       ShutdownHandler
-	proxyResourceHandler       *handler.StreamWatchProxyHandler
-	isInitialized              bool
+	teamMap                       cache.Cache
+	cacheManager                  agentcache.Manager
+	apiValidator                  APIValidator
+	apiValidatorLock              sync.Mutex
+	apiValidatorJobID             string
+	configChangeHandler           ConfigChangeHandler
+	agentResourceChangeHandler    ConfigChangeHandler
+	customUnitMetricServerManager customunit.CustomUnitMetricServerManager
+	agentShutdownHandler          ShutdownHandler
+	proxyResourceHandler          *handler.StreamWatchProxyHandler
+	isInitialized                 bool
 
 	provisioner          provisioning.Provisioning
 	streamer             *stream.StreamerClient
@@ -170,16 +169,10 @@ func InitializeWithAgentFeatures(centralCfg config.CentralConfig, agentFeaturesC
 
 	// call the metric services.
 	metricServicesConfigs := agentFeaturesCfg.GetMetricServicesConfigs()
-	agent.customUnitQuotaHandler = customunithandler.NewCustomUnitQuotaHandler(metricServicesConfigs)
-	// iterate over each metric service config
-	for _, config := range metricServicesConfigs {
-		ctx, ctxCancel := context.WithCancel(context.Background())
-		// Initialize custom units client
-		factory := customunit.NewCustomMetricReportingClientFactory(config.URL, agent.cacheManager)
-		client, _ := factory(ctx, ctxCancel)
+	agent.customUnitMetricServerManager = customunit.NewCustomUnitMetricServerManager(metricServicesConfigs, agent.cacheManager)
+	ctx, ctxCancel := context.WithCancel(context.Background())
+	agent.customUnitMetricServerManager.HandleMetricReporting(ctx, ctxCancel)
 
-		client.MetricReporting()
-	}
 	if !agent.isInitialized {
 		err = handleInitialization()
 		if err != nil {

diff --git a/pkg/agent/handler/accessrequest.go b/pkg/agent/handler/accessrequest.go
@@ -9,7 +9,6 @@ import (
 	management "github.com/Axway/agent-sdk/pkg/apic/apiserver/models/management/v1alpha1"
 	defs "github.com/Axway/agent-sdk/pkg/apic/definitions"
 	prov "github.com/Axway/agent-sdk/pkg/apic/provisioning"
-	handler "github.com/Axway/agent-sdk/pkg/customunit/handler"
 	"github.com/Axway/agent-sdk/pkg/util"
 	"github.com/Axway/agent-sdk/pkg/watchmanager/proto"
 )
@@ -24,24 +23,27 @@ type arProvisioner interface {
 	AccessRequestProvision(accessRequest prov.AccessRequest) (status prov.RequestStatus, data prov.AccessData)
 	AccessRequestDeprovision(accessRequest prov.AccessRequest) (status prov.RequestStatus)
 }
+type customUnitMetricServerManager interface {
+	HandleQuotaEnforcement(context.Context, context.CancelFunc, *management.AccessRequest, *management.ManagedApplication) error
+}
 
 type accessRequestHandler struct {
 	marketplaceHandler
-	prov                   arProvisioner
-	cache                  agentcache.Manager
-	client                 client
-	encryptSchema          encryptSchemaFunc
-	customUnitQuotaHandler handler.CustomUnitQuotaHandler
+	prov                          arProvisioner
+	cache                         agentcache.Manager
+	client                        client
+	encryptSchema                 encryptSchemaFunc
+	customUnitMetricServerManager customUnitMetricServerManager
 }
 
 // NewAccessRequestHandler creates a Handler for Access Requests
-func NewAccessRequestHandler(prov arProvisioner, cache agentcache.Manager, client client, customUnitQuotaHandler handler.CustomUnitQuotaHandler) Handler {
+func NewAccessRequestHandler(prov arProvisioner, cache agentcache.Manager, client client, customUnitMetricServerManager customUnitMetricServerManager) Handler {
 	return &accessRequestHandler{
-		prov:                   prov,
-		cache:                  cache,
-		client:                 client,
-		encryptSchema:          encryptSchema,
-		customUnitQuotaHandler: customUnitQuotaHandler,
+		prov:                          prov,
+		cache:                         cache,
+		client:                        client,
+		encryptSchema:                 encryptSchema,
+		customUnitMetricServerManager: customUnitMetricServerManager,
 	}
 }
 
@@ -137,15 +139,13 @@ func (h *accessRequestHandler) onPending(ctx context.Context, ar *management.Acc
 	status, accessData := h.prov.AccessRequestProvision(req)
 
 	if status.GetStatus() == prov.Success && len(ar.Spec.AdditionalQuotas) > 0 {
-		errMessage, err := h.customUnitQuotaHandler.Handle(ctx, ar, app)
+		ctx, cancelCtx := context.WithCancel(ctx)
+		err := h.customUnitMetricServerManager.HandleQuotaEnforcement(ctx, cancelCtx, ar, app)
 
 		if err != nil {
-			h.onError(ctx, ar, err)
-		}
-
-		if errMessage != "" {
+			// h.onError(ctx, ar, err)
 			status = prov.NewRequestStatusBuilder().
-				SetMessage(errMessage).
+				SetMessage(err.Error()).
 				SetCurrentStatusReasons(ar.Status.Reasons).
 				Failed()
 		}

diff --git a/pkg/agent/handler/accessrequest_test.go b/pkg/agent/handler/accessrequest_test.go
@@ -12,7 +12,7 @@ import (
 	prov "github.com/Axway/agent-sdk/pkg/apic/provisioning"
 	"github.com/Axway/agent-sdk/pkg/apic/provisioning/mock"
 	"github.com/Axway/agent-sdk/pkg/config"
-	handler "github.com/Axway/agent-sdk/pkg/customunit/handler"
+	"github.com/Axway/agent-sdk/pkg/customunit"
 	"github.com/Axway/agent-sdk/pkg/util"
 	"github.com/Axway/agent-sdk/pkg/watchmanager/proto"
 	"github.com/stretchr/testify/assert"
@@ -176,8 +176,8 @@ func TestAccessRequestHandler(t *testing.T) {
 				c.isDeleting = true
 			}
 			af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-			customUnitQuotaHandler := handler.NewCustomUnitQuotaHandler(af)
-			handler := NewAccessRequestHandler(arp, cm, c, customUnitQuotaHandler)
+			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+			handler := NewAccessRequestHandler(arp, cm, c, customUnitMetricServerManager)
 			v := handler.(*accessRequestHandler)
 			v.encryptSchema = func(_, _ map[string]interface{}, _, _, _ string) (map[string]interface{}, error) {
 				return map[string]interface{}{}, nil
@@ -251,8 +251,8 @@ func TestAccessRequestHandler_deleting(t *testing.T) {
 				t:              t,
 			}
 			af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-			customUnitQuotaHandler := handler.NewCustomUnitQuotaHandler(af)
-			handler := NewAccessRequestHandler(arp, cm, c, customUnitQuotaHandler)
+			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+			handler := NewAccessRequestHandler(arp, cm, c, customUnitMetricServerManager)
 
 			ri, _ := ar.AsInstance()
 
@@ -276,8 +276,8 @@ func TestAccessRequestHandler_wrong_kind(t *testing.T) {
 	}
 	ar := &mockARProvision{}
 	af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-	customUnitQuotaHandler := handler.NewCustomUnitQuotaHandler(af)
-	handler := NewAccessRequestHandler(ar, cm, c, customUnitQuotaHandler)
+	customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+	handler := NewAccessRequestHandler(ar, cm, c, customUnitMetricServerManager)
 	ri := &v1.ResourceInstance{
 		ResourceMeta: v1.ResourceMeta{
 			GroupVersionKind: management.EnvironmentGVK(),

diff --git a/pkg/agent/provisioning.go b/pkg/agent/provisioning.go
@@ -537,7 +537,7 @@ func RegisterProvisioner(provisioner provisioning.Provisioning) {
 	if agent.cfg.GetAgentType() == config.DiscoveryAgent || agent.cfg.GetAgentType() == config.GovernanceAgent {
 		agent.proxyResourceHandler.RegisterTargetHandler(
 			"accessrequesthandler",
-			handler.NewAccessRequestHandler(agent.provisioner, agent.cacheManager, agent.apicClient, agent.customUnitQuotaHandler),
+			handler.NewAccessRequestHandler(agent.provisioner, agent.cacheManager, agent.apicClient, agent.customUnitMetricServerManager),
 		)
 		agent.proxyResourceHandler.RegisterTargetHandler(
 			"managedappHandler",