APIGOV-28984 Updates

Axway · Nov 5, 2024 · 77e88a2 · 77e88a2
1 parent 0b424b3
commit 77e88a2
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 26 deletions.
diff --git a/pkg/agent/agent.go b/pkg/agent/agent.go
@@ -69,7 +69,7 @@ type agentData struct {
 	apiValidatorJobID             string
 	configChangeHandler           ConfigChangeHandler
 	agentResourceChangeHandler    ConfigChangeHandler
-	customUnitMetricServerManager customunit.CustomUnitMetricServerManager
+	customUnitMetricServerManager *customunit.CustomUnitMetricServerManager
 	agentShutdownHandler          ShutdownHandler
 	proxyResourceHandler          *handler.StreamWatchProxyHandler
 	isInitialized                 bool
@@ -169,7 +169,7 @@ func InitializeWithAgentFeatures(centralCfg config.CentralConfig, agentFeaturesC
 
 	// call the metric services.
 	metricServicesConfigs := agentFeaturesCfg.GetMetricServicesConfigs()
-	agent.customUnitMetricServerManager = customunit.NewCustomUnitMetricServerManager(metricServicesConfigs, agent.cacheManager)
+	agent.customUnitMetricServerManager = customunit.NewCustomUnitMetricServerManager(metricServicesConfigs, agent.cacheManager, centralCfg.GetAgentType())
 	ctx, ctxCancel := context.WithCancel(context.Background())
 	agent.customUnitMetricServerManager.HandleMetricReporting(ctx, ctxCancel)
 

diff --git a/pkg/agent/handler/accessrequest_test.go b/pkg/agent/handler/accessrequest_test.go
@@ -176,7 +176,7 @@ func TestAccessRequestHandler(t *testing.T) {
 				c.isDeleting = true
 			}
 			af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm, config.DiscoveryAgent)
 			handler := NewAccessRequestHandler(arp, cm, c, customUnitMetricServerManager)
 			v := handler.(*accessRequestHandler)
 			v.encryptSchema = func(_, _ map[string]interface{}, _, _, _ string) (map[string]interface{}, error) {
@@ -251,7 +251,7 @@ func TestAccessRequestHandler_deleting(t *testing.T) {
 				t:              t,
 			}
 			af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+			customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm, config.DiscoveryAgent)
 			handler := NewAccessRequestHandler(arp, cm, c, customUnitMetricServerManager)
 
 			ri, _ := ar.AsInstance()
@@ -276,7 +276,7 @@ func TestAccessRequestHandler_wrong_kind(t *testing.T) {
 	}
 	ar := &mockARProvision{}
 	af := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
-	customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm)
+	customUnitMetricServerManager := customunit.NewCustomUnitMetricServerManager(af, cm, config.DiscoveryAgent)
 	handler := NewAccessRequestHandler(ar, cm, c, customUnitMetricServerManager)
 	ri := &v1.ResourceInstance{
 		ResourceMeta: v1.ResourceMeta{

diff --git a/pkg/customunit/client.go b/pkg/customunit/client.go
@@ -214,7 +214,7 @@ func (c *customUnitClient) APIServiceLookup(apiServiceLookup *cu.APIServiceLooku
 	case cu.APIServiceLookupType_ExternalAPIID:
 		apiSvc = c.cache.GetAPIServiceWithAPIID(apiSvcValue)
 	case cu.APIServiceLookupType_ServiceID:
-		apiSvc = c.cache.GetAPIServiceWithPrimaryKey(apiSvcValue)
+		apiSvc = c.cache.GetAPIServiceWithAPIID(apiSvcValue)
 	case cu.APIServiceLookupType_ServiceName:
 		apiSvc = c.cache.GetAPIServiceWithName(apiSvcValue)
 	}
@@ -263,10 +263,14 @@ func (c *customUnitClient) ManagedApplicationLookup(appLookup *cu.AppLookup) (*m
 	if managedApp == nil {
 		return nil, nil
 	}
+	consumerOrgID := ""
+	if managedApp.Owner != nil {
+		consumerOrgID = managedApp.Owner.ID
+	}
 	return &models.AppDetails{
 		ID:            managedApp.Metadata.ID,
 		Name:          managedApp.Name,
-		ConsumerOrgID: managedApp.Owner.ID,
+		ConsumerOrgID: consumerOrgID,
 	}, nil
 }
 

diff --git a/pkg/customunit/manager.go b/pkg/customunit/manager.go
@@ -15,24 +15,21 @@ import (
 	"github.com/Axway/agent-sdk/pkg/util"
 )
 
-type customUnitMetricServerManager struct {
-	configs []config.MetricServiceConfiguration
-	cache   cache.Manager
+type CustomUnitMetricServerManager struct {
+	configs   []config.MetricServiceConfiguration
+	cache     cache.Manager
+	agentType config.AgentType
 }
 
-type CustomUnitMetricServerManager interface {
-	HandleQuotaEnforcement(context.Context, context.CancelFunc, *management.AccessRequest, *management.ManagedApplication) error
-	HandleMetricReporting(context.Context, context.CancelFunc)
-}
-
-func NewCustomUnitMetricServerManager(configs []config.MetricServiceConfiguration, cache cache.Manager) CustomUnitMetricServerManager {
-	return &customUnitMetricServerManager{
-		configs: configs,
-		cache:   cache,
+func NewCustomUnitMetricServerManager(configs []config.MetricServiceConfiguration, cache cache.Manager, agentType config.AgentType) *CustomUnitMetricServerManager {
+	return &CustomUnitMetricServerManager{
+		configs:   configs,
+		cache:     cache,
+		agentType: agentType,
 	}
 }
 
-func (h *customUnitMetricServerManager) HandleQuotaEnforcement(ctx context.Context, cancelCtx context.CancelFunc, ar *management.AccessRequest, app *management.ManagedApplication) error {
+func (h *CustomUnitMetricServerManager) HandleQuotaEnforcement(ctx context.Context, cancelCtx context.CancelFunc, ar *management.AccessRequest, app *management.ManagedApplication) error {
 	// Build quota info
 	quotaInfo, err := h.buildQuotaInfo(ctx, ar, app)
 	if err != nil {
@@ -59,7 +56,7 @@ func (h *customUnitMetricServerManager) HandleQuotaEnforcement(ctx context.Conte
 	return nil
 }
 
-func (h *customUnitMetricServerManager) buildQuotaInfo(ctx context.Context, ar *management.AccessRequest, app *management.ManagedApplication) (*customunits.QuotaInfo, error) {
+func (h *CustomUnitMetricServerManager) buildQuotaInfo(ctx context.Context, ar *management.AccessRequest, app *management.ManagedApplication) (*customunits.QuotaInfo, error) {
 	unitRef, count := h.getQuotaInfo(ar)
 	if unitRef == "" {
 		return nil, nil
@@ -108,7 +105,7 @@ type reference struct {
 	Unit string `json:"unit"`
 }
 
-func (h *customUnitMetricServerManager) getQuotaInfo(ar *management.AccessRequest) (string, int) {
+func (h *CustomUnitMetricServerManager) getQuotaInfo(ar *management.AccessRequest) (string, int) {
 	index := 0
 	if len(ar.Spec.AdditionalQuotas) < index+1 {
 		return "", 0
@@ -126,7 +123,7 @@ func (h *customUnitMetricServerManager) getQuotaInfo(ar *management.AccessReques
 	return "", 0
 }
 
-func (h *customUnitMetricServerManager) getServiceInstance(_ context.Context, ar *management.AccessRequest) (*apiv1.ResourceInstance, error) {
+func (h *CustomUnitMetricServerManager) getServiceInstance(_ context.Context, ar *management.AccessRequest) (*apiv1.ResourceInstance, error) {
 	instRef := ar.GetReferenceByGVK(management.APIServiceInstanceGVK())
 	instID := instRef.ID
 	instance, err := h.cache.GetAPIServiceInstanceByID(instID)
@@ -136,7 +133,10 @@ func (h *customUnitMetricServerManager) getServiceInstance(_ context.Context, ar
 	return instance, nil
 }
 
-func (m *customUnitMetricServerManager) HandleMetricReporting(ctx context.Context, cancelCtx context.CancelFunc) {
+func (m *CustomUnitMetricServerManager) HandleMetricReporting(ctx context.Context, cancelCtx context.CancelFunc) {
+	if m.agentType != config.TraceabilityAgent {
+		return
+	}
 	// iterate over each metric service config
 	for _, config := range m.configs {
 		// Initialize custom units client

diff --git a/pkg/customunit/manager_test.go b/pkg/customunit/manager_test.go
@@ -20,7 +20,7 @@ const managedAppRefName = "managed-app-name"
 func Test_NewCustomUnitMetricServerManager(t *testing.T) {
 	metricServicesConfigs := config.NewAgentFeaturesConfiguration().GetMetricServicesConfigs()
 	cm := agentcache.NewAgentCacheManager(&config.CentralConfiguration{}, false)
-	handler := NewCustomUnitMetricServerManager(metricServicesConfigs, cm)
+	handler := NewCustomUnitMetricServerManager(metricServicesConfigs, cm, config.DiscoveryAgent)
 
 	assert.NotNil(t, handler)
 }
@@ -91,7 +91,7 @@ func Test_HandleQuotaEnforcementInfo(t *testing.T) {
 		},
 	}
 
-	manager := NewCustomUnitMetricServerManager(metricServicesConfigs, cm)
+	manager := NewCustomUnitMetricServerManager(metricServicesConfigs, cm, config.DiscoveryAgent)
 	ctx, cancelCtx := context.WithCancel(context.Background())
 	err := manager.HandleQuotaEnforcement(ctx, cancelCtx, accessReq, managedAppForTest)