Add company selection handling in customer forms with unit tests

- Introduced new functionality in `form_companies.go` to handle company selection from various input formats, including legacy `company_ids` and company summary objects.
- Implemented `UnmarshalJSON` methods for `CreateCustomerForm`, `UpdateCustomerForm`, and `AssignCompaniesForm` to support flexible company data parsing.
- Added unit tests in `form_companies_test.go` to validate the correct unmarshalling of company IDs and summaries, ensuring robust handling of different input scenarios.
- Enhanced the `GetByCompanies` method in the customer repository to directly use string company IDs, improving data retrieval efficiency.

This update significantly improves the handling of company data in customer forms, ensuring accurate processing and validation of company selections, while expanding test coverage for these functionalities.

internal/customer/form_companies.go

@@ -0,0 +1,152 @@
+package customer
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+type companySelectionObject struct {
+	ID string `json:"id"`
+}
+
+// decodeCompanySelection parses company IDs sent as strings, legacy company_ids,
+// or company summary objects returned by GET customer.
+func decodeCompanySelection(raw json.RawMessage) ([]string, error) {
+	if len(raw) == 0 {
+		return nil, fmt.Errorf("empty company selection")
+	}
+	if string(raw) == "null" {
+		return []string{}, nil
+	}
+
+	var ids []string
+	if err := json.Unmarshal(raw, &ids); err == nil {
+		return ids, nil
+	}
+
+	var objects []companySelectionObject
+	if err := json.Unmarshal(raw, &objects); err == nil {
+		result := make([]string, 0, len(objects))
+		for _, object := range objects {
+			if id := strings.TrimSpace(object.ID); id != "" {
+				result = append(result, id)
+			}
+		}
+		return result, nil
+	}
+
+	return nil, fmt.Errorf("companies must be an array of IDs or objects with id")
+}
+
+func mergeCompanySelectionFields(raw map[string]json.RawMessage) (json.RawMessage, bool, error) {
+	if value, ok := raw["companies"]; ok {
+		return value, true, nil
+	}
+	if value, ok := raw["company_ids"]; ok {
+		return value, true, nil
+	}
+	return nil, false, nil
+}
+
+func unmarshalFormWithoutCompanyFields(data []byte, raw map[string]json.RawMessage, target any) error {
+	delete(raw, "companies")
+	delete(raw, "company_ids")
+
+	trimmed, err := json.Marshal(raw)
+	if err != nil {
+		return err
+	}
+
+	return json.Unmarshal(trimmed, target)
+}
+
+// UnmarshalJSON accepts companies, legacy company_ids, and company summary objects.
+func (f *CreateCustomerForm) UnmarshalJSON(data []byte) error {
+	var raw map[string]json.RawMessage
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	selectionRaw, present, err := mergeCompanySelectionFields(raw)
+	if err != nil {
+		return err
+	}
+
+	type alias CreateCustomerForm
+	if err := unmarshalFormWithoutCompanyFields(data, raw, (*alias)(f)); err != nil {
+		return err
+	}
+
+	if !present {
+		return nil
+	}
+
+	companies, err := decodeCompanySelection(selectionRaw)
+	if err != nil {
+		return err
+	}
+	f.Companies = companies
+
+	return nil
+}
+
+// UnmarshalJSON accepts companies, legacy company_ids, and company summary objects.
+func (f *UpdateCustomerForm) UnmarshalJSON(data []byte) error {
+	var raw map[string]json.RawMessage
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	selectionRaw, present, err := mergeCompanySelectionFields(raw)
+	if err != nil {
+		return err
+	}
+
+	type alias UpdateCustomerForm
+	if err := unmarshalFormWithoutCompanyFields(data, raw, (*alias)(f)); err != nil {
+		return err
+	}
+
+	if !present {
+		return nil
+	}
+
+	companies, err := decodeCompanySelection(selectionRaw)
+	if err != nil {
+		return err
+	}
+	f.Companies = &companies
+
+	return nil
+}
+
+// UnmarshalJSON accepts companies, legacy company_ids, and company summary objects.
+func (f *AssignCompaniesForm) UnmarshalJSON(data []byte) error {
+	var raw map[string]json.RawMessage
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return err
+	}
+
+	selectionRaw, present, err := mergeCompanySelectionFields(raw)
+	if err != nil {
+		return err
+	}
+
+	type alias AssignCompaniesForm
+	if err := unmarshalFormWithoutCompanyFields(data, raw, (*alias)(f)); err != nil {
+		return err
+	}
+
+	if !present {
+		return nil
+	}
+
+	companies, err := decodeCompanySelection(selectionRaw)
+	if err != nil {
+		return err
+	}
+	f.Companies = companies
+
+	return nil
+}

internal/customer/form_companies_test.go

@@ -0,0 +1,92 @@
+package customer
+
+import (
+	"encoding/json"
+	"testing"
+)
+
+func TestCreateCustomerForm_UnmarshalJSON_CompanyIDs(t *testing.T) {
+	payload := `{
+		"type": "individual",
+		"role": "analyst",
+		"username": "testuser",
+		"email": "test@example.com",
+		"password": "Test!1234",
+		"company_ids": ["507f1f77bcf86cd799439011", "507f1f77bcf86cd799439012"]
+	}`
+
+	var form CreateCustomerForm
+	if err := json.Unmarshal([]byte(payload), &form); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	want := []string{"507f1f77bcf86cd799439011", "507f1f77bcf86cd799439012"}
+	if len(form.Companies) != len(want) {
+		t.Fatalf("Companies len = %d, want %d", len(form.Companies), len(want))
+	}
+	for i, id := range want {
+		if form.Companies[i] != id {
+			t.Fatalf("Companies[%d] = %q, want %q", i, form.Companies[i], id)
+		}
+	}
+}
+
+func TestUpdateCustomerForm_UnmarshalJSON_CompanySummaries(t *testing.T) {
+	payload := `{
+		"type": "individual",
+		"role": "analyst",
+		"username": "testuser",
+		"email": "test@example.com",
+		"companies": [
+			{"id": "507f1f77bcf86cd799439011", "name": "Acme"},
+			{"id": "507f1f77bcf86cd799439012", "name": "Globex"}
+		]
+	}`
+
+	var form UpdateCustomerForm
+	if err := json.Unmarshal([]byte(payload), &form); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if form.Companies == nil {
+		t.Fatal("Companies should be set")
+	}
+
+	want := []string{"507f1f77bcf86cd799439011", "507f1f77bcf86cd799439012"}
+	if len(*form.Companies) != len(want) {
+		t.Fatalf("Companies len = %d, want %d", len(*form.Companies), len(want))
+	}
+	for i, id := range want {
+		if (*form.Companies)[i] != id {
+			t.Fatalf("Companies[%d] = %q, want %q", i, (*form.Companies)[i], id)
+		}
+	}
+}
+
+func TestUpdateCustomerForm_UnmarshalJSON_OmittedCompanies(t *testing.T) {
+	payload := `{
+		"type": "individual",
+		"role": "analyst",
+		"username": "testuser",
+		"email": "test@example.com"
+	}`
+
+	var form UpdateCustomerForm
+	if err := json.Unmarshal([]byte(payload), &form); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if form.Companies != nil {
+		t.Fatalf("Companies = %v, want nil when omitted", form.Companies)
+	}
+}
+
+func TestAssignCompaniesForm_UnmarshalJSON_LegacyField(t *testing.T) {
+	payload := `{"company_ids": ["507f1f77bcf86cd799439011"]}`
+
+	var form AssignCompaniesForm
+	if err := json.Unmarshal([]byte(payload), &form); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if len(form.Companies) != 1 || form.Companies[0] != "507f1f77bcf86cd799439011" {
+		t.Fatalf("Companies = %v", form.Companies)
+	}
+}

internal/customer/repository.go

@@ -249,16 +249,11 @@ func (r *customerRepository) GetByRole(ctx context.Context, role CustomerRole) (
 // GetByCompanies retrieves customers by companies
 
 func (r *customerRepository) GetByCompanies(ctx context.Context, companies []string) ([]Customer, error) {
-	companyIDs := make([]bson.ObjectID, 0)
-	for _, company := range companies {
-		id, err := bson.ObjectIDFromHex(company)
-		if err != nil {
-			return nil, err
-		}
-		companyIDs = append(companyIDs, id)
+	if len(companies) == 0 {
+		return []Customer{}, nil
 	}
 
-	customers, err := r.ormRepo.FindAll(ctx, bson.M{"companies": bson.M{"$in": companyIDs}}, orm.NewPaginationBuilder().Build())
+	customers, err := r.ormRepo.FindAll(ctx, bson.M{"companies": bson.M{"$in": companies}}, orm.NewPaginationBuilder().Build())
 	if err != nil {
 		return nil, err
 	}
@@ -274,10 +269,10 @@ func (r *customerRepository) GetByCompanyID(ctx context.Context, companyID strin
 		SortBy(pagination.SortBy, pagination.SortOrder).
 		Build()
 
-	// Filter by company ID and active status
+	// Filter by assigned company ID and active status
 	filter := bson.M{
-		"company_id": companyID,
-		"status":     bson.M{"$ne": CustomerStatusInactive},
+		"companies": companyID,
+		"status":    bson.M{"$ne": CustomerStatusInactive},
 	}
 
 	// Use ORM to find customers