[usage] Return CostCenter from SetCostCenter udpates - implementation

Author: easyCZ

components/usage/pkg/apiv1/usage.go

@@ -185,15 +185,19 @@ func (s *UsageService) GetCostCenter(ctx context.Context, in *v1.GetCostCenterRe
 		return nil, err
 	}
 	return &v1.GetCostCenterResponse{
-		CostCenter: &v1.CostCenter{
-			AttributionId:   string(result.ID),
-			SpendingLimit:   result.SpendingLimit,
-			BillingStrategy: convertBillingStrategyToAPI(result.BillingStrategy),
-			NextBillingTime: timestamppb.New(result.NextBillingTime.Time()),
-		},
+		CostCenter: dbCostCenterToAPI(result),
 	}, nil
 }
 
+func dbCostCenterToAPI(c db.CostCenter) *v1.CostCenter {
+	return &v1.CostCenter{
+		AttributionId:   string(c.ID),
+		SpendingLimit:   c.SpendingLimit,
+		BillingStrategy: convertBillingStrategyToAPI(c.BillingStrategy),
+		NextBillingTime: timestamppb.New(c.NextBillingTime.Time()),
+	}
+}
+
 func convertBillingStrategyToDB(in v1.CostCenter_BillingStrategy) db.BillingStrategy {
 	if in == v1.CostCenter_BILLING_STRATEGY_STRIPE {
 		return db.CostCenter_Stripe
@@ -223,11 +227,13 @@ func (s *UsageService) SetCostCenter(ctx context.Context, in *v1.SetCostCenterRe
 		SpendingLimit:   in.CostCenter.SpendingLimit,
 		BillingStrategy: convertBillingStrategyToDB(in.CostCenter.BillingStrategy),
 	}
-	_, err = s.costCenterManager.UpdateCostCenter(ctx, costCenter)
+	result, err := s.costCenterManager.UpdateCostCenter(ctx, costCenter)
 	if err != nil {
 		return nil, err
 	}
-	return &v1.SetCostCenterResponse{}, nil
+	return &v1.SetCostCenterResponse{
+		CostCenter: dbCostCenterToAPI(result),
+	}, nil
 }
 
 func (s *UsageService) ReconcileUsage(ctx context.Context, req *v1.ReconcileUsageRequest) (*v1.ReconcileUsageResponse, error) {

components/usage/pkg/apiv1/usage_test.go

@@ -251,14 +251,10 @@ func TestGetAndSetCostCenter(t *testing.T) {
 	service := newUsageService(t, conn)
 
 	for _, costCenter := range costCenterUpdates {
-		_, err := service.SetCostCenter(context.Background(), &v1.SetCostCenterRequest{
+		retrieved, err := service.SetCostCenter(context.Background(), &v1.SetCostCenterRequest{
 			CostCenter: costCenter,
 		})
 		require.NoError(t, err)
-		retrieved, err := service.GetCostCenter(context.Background(), &v1.GetCostCenterRequest{
-			AttributionId: costCenter.AttributionId,
-		})
-		require.NoError(t, err)
 
 		require.Equal(t, costCenter.SpendingLimit, retrieved.CostCenter.SpendingLimit)
 		require.Equal(t, costCenter.BillingStrategy, retrieved.CostCenter.BillingStrategy)