fix: update hook tests and rename inputs

Author: GAlexIHU

openmeter/customer/adapter/customer.go

@@ -464,14 +464,14 @@ func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input custo
 				customerdb.Or(
 					// We lookup the customer by subject key in the subjects table
 					customerdb.HasSubjectsWith(
-						customersubjectsdb.SubjectKey(input.SubjectKey),
+						customersubjectsdb.SubjectKey(input.Key),
 						customersubjectsdb.Or(
 							customersubjectsdb.DeletedAtIsNil(),
 							customersubjectsdb.DeletedAtGT(now),
 						),
 					),
 					// Or else we lookup the customer by key in the customers table
-					customerdb.Key(input.SubjectKey),
+					customerdb.Key(input.Key),
 				),
 			).
 			Where(customerdb.DeletedAtIsNil())
@@ -484,7 +484,7 @@ func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input custo
 		if err != nil {
 			if entdb.IsNotFound(err) {
 				return nil, models.NewGenericNotFoundError(
-					fmt.Errorf("customer with subject key %s not found in %s namespace", input.SubjectKey, input.Namespace),
+					fmt.Errorf("customer with subject key %s not found in %s namespace", input.Key, input.Namespace),
 				)
 			}
 

openmeter/customer/customer.go

@@ -212,8 +212,10 @@ func (c CustomerUsageAttribution) GetSubjectKey() (string, error) {
 
 // GetCustomerByUsageAttributionInput represents the input for the GetCustomerByUsageAttribution method
 type GetCustomerByUsageAttributionInput struct {
-	Namespace  string
-	SubjectKey string
+	Namespace string
+
+	// The key of either the customer or one of its subjects
+	Key string
 
 	// Expand
 	Expands Expands
@@ -224,7 +226,7 @@ func (i GetCustomerByUsageAttributionInput) Validate() error {
 		return models.NewGenericValidationError(errors.New("namespace is required"))
 	}
 
-	if i.SubjectKey == "" {
+	if i.Key == "" {
 		return models.NewGenericValidationError(errors.New("subject key is required"))
 	}
 

openmeter/customer/service/hooks/subjectcustomer.go

@@ -230,8 +230,8 @@ var ErrCustomerKeyConflict = errors.New("customer key conflict")
 func (p CustomerProvisioner) getCustomerForSubject(ctx context.Context, sub *subject.Subject) (*customer.Customer, error) {
 	// Try to find Customer for Subject by usage attribution
 	cus, err := p.customer.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  sub.Namespace,
-		SubjectKey: sub.Key,
+		Namespace: sub.Namespace,
+		Key:       sub.Key,
 	})
 	if err != nil && !models.IsGenericNotFoundError(err) {
 		return nil, err

openmeter/customer/service/hooks/subjectcustomer_test.go

@@ -150,7 +150,11 @@ func TestCustomerProvisioner_EnsureCustomer(t *testing.T) {
 			require.NoError(t, err, "creating subject should not fail")
 			assert.NotNilf(t, sub, "subject must not be nil")
 
-			cus, err := env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{
+			cusForSubject, err := provisioner.EnsureCustomer(ctx, &sub)
+			require.NoError(t, err, "provisioning customer should not fail")
+			assert.NotNilf(t, cusForSubject, "customer must not be nil")
+
+			_, err = env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{
 				Namespace: namespace,
 				CustomerMutate: customer.CustomerMutate{
 					Key:         lo.ToPtr(sub.Key),
@@ -170,22 +174,8 @@ func TestCustomerProvisioner_EnsureCustomer(t *testing.T) {
 					Annotation: nil,
 				},
 			})
-			require.NoError(t, err, "creating customer should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
-
-			cus, err = provisioner.EnsureCustomer(ctx, &sub)
-			require.NoError(t, err, "provisioning customer should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
 
-			cus, err = env.CustomerService.GetCustomer(ctx, customer.GetCustomerInput{
-				CustomerID: &customer.CustomerID{
-					Namespace: cus.Namespace,
-					ID:        cus.ID,
-				},
-			})
-			require.NoErrorf(t, err, "getting customer for subject should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
-			AssertSubjectCustomerEqual(t, &sub, cus)
+			require.True(t, models.IsGenericConflictError(err), "creating customer should fail with conflict")
 		})
 
 		t.Run("CustomerKeyMismatch", func(t *testing.T) {

openmeter/customer/service/hooks/subjectvalidator.go

@@ -28,8 +28,8 @@ type subjectValidatorHook struct {
 
 func (s subjectValidatorHook) PreDelete(ctx context.Context, sub *subject.Subject) error {
 	cus, err := s.customer.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  sub.Namespace,
-		SubjectKey: sub.Key,
+		Namespace: sub.Namespace,
+		Key:       sub.Key,
 	})
 	if err != nil {
 		if models.IsGenericNotFoundError(err) {

openmeter/customer/service/service_test.go

@@ -86,8 +86,8 @@ func Test_CustomerService(t *testing.T) {
 
 				t.Run("ByUsageAttribution", func(t *testing.T) {
 					cusByUsage, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-						Namespace:  cus.Namespace,
-						SubjectKey: cus.UsageAttribution.SubjectKeys[0],
+						Namespace: cus.Namespace,
+						Key:       cus.UsageAttribution.SubjectKeys[0],
 					})
 					require.NoError(t, err, "getting customer usage attribution should not fail")
 					assert.NotNilf(t, cusByUsage, "customer must not be nil")
@@ -128,8 +128,8 @@ func Test_CustomerService(t *testing.T) {
 
 				t.Run("ByUsageAttribution", func(t *testing.T) {
 					cusByUsage, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-						Namespace:  cus.Namespace,
-						SubjectKey: subjectKeys[1],
+						Namespace: cus.Namespace,
+						Key:       subjectKeys[1],
 					})
 					require.NoError(t, err, "getting customer usage attribution should not fail")
 					assert.NotNilf(t, cusByUsage, "customer must not be nil")

openmeter/entitlement/driver/entitlement.go

@@ -569,8 +569,8 @@ func (h *entitlementHandler) resolveCustomerFromSubject(ctx context.Context, nam
 	}
 
 	cust, err := h.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  namespace,
-		SubjectKey: subj.Key,
+		Namespace: namespace,
+		Key:       subj.Key,
 	})
 	if err != nil {
 		return nil, err

openmeter/entitlement/driver/metered.go

@@ -395,8 +395,8 @@ func (h *meteredEntitlementHandler) resolveCustomerFromSubject(ctx context.Conte
 	}
 
 	cust, err := h.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  namespace,
-		SubjectKey: subj.Key,
+		Namespace: namespace,
+		Key:       subj.Key,
 	})
 	if err != nil {
 		return nil, err

openmeter/meterevent/adapter/event.go

@@ -277,8 +277,8 @@ func (a *adapter) enrichEventsWithCustomerID(ctx context.Context, namespace stri
 		// FIXME: do this in a batches to avoid hitting the database for each event
 		// Get the customer by usage attribution subject key
 		cust, err := a.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-			Namespace:  namespace,
-			SubjectKey: event.Subject,
+			Namespace: namespace,
+			Key:       event.Subject,
 		})
 		if err != nil {
 			if models.IsGenericNotFoundError(err) {

test/customer/customer.go

@@ -700,8 +700,8 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 
 	// Get the customer by usage attribution
 	cus, err := service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: TestSubjectKeys[0],
+		Namespace: s.namespace,
+		Key:       TestSubjectKeys[0],
 	})
 
 	require.NoError(t, err, "Fetching customer must not return error")
@@ -711,8 +711,8 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 
 	// Get the customer by key
 	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: TestKey,
+		Namespace: s.namespace,
+		Key:       TestKey,
 	})
 
 	require.NoError(t, err, "Fetching customer must not return error")
@@ -722,8 +722,8 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 
 	// Get the customer by key
 	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: TestKey,
+		Namespace: s.namespace,
+		Key:       TestKey,
 	})
 
 	require.NoError(t, err, "Fetching customer must not return error")
@@ -733,8 +733,8 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 
 	// Get the customer by usage attribution with a non-existent subject key
 	_, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: "non-existent-subject-key",
+		Namespace: s.namespace,
+		Key:       "non-existent-subject-key",
 	})
 
 	require.True(t, models.IsGenericNotFoundError(err), "Fetching customer with non-existent subject key must return not found error")