[FSSDK-11133] Add cmab to project config (#399)

* add cmab to project config

* reformat

* add attribute key/id mapping to config

* add traffic allocation inside cmab

* add test case to test mapCmab

* change cmab traffic allocation type to int

* add nil test

Author: Mat001

pkg/config/datafileprojectconfig/config.go

@@ -1,5 +1,5 @@
 /****************************************************************************
- * Copyright 2019-2022, Optimizely, Inc. and contributors                   *
+ * Copyright 2019-2025, Optimizely, Inc. and contributors                   *
  *                                                                          *
  * Licensed under the Apache License, Version 2.0 (the "License");          *
  * you may not use this file except in compliance with the License.         *
@@ -41,8 +41,10 @@ type DatafileProjectConfig struct {
 	experimentKeyToIDMap map[string]string
 	audienceMap          map[string]entities.Audience
 	attributeMap         map[string]entities.Attribute
+	attributeKeyMap      map[string]entities.Attribute
 	eventMap             map[string]entities.Event
 	attributeKeyToIDMap  map[string]string
+	attributeIDToKeyMap  map[string]string
 	experimentMap        map[string]entities.Experiment
 	featureMap           map[string]entities.Feature
 	groupMap             map[string]entities.Group
@@ -107,6 +109,24 @@ func (c DatafileProjectConfig) GetAttributeID(key string) string {
 	return c.attributeKeyToIDMap[key]
 }
 
+// GetAttributeByKey returns the attribute with the given key
+func (c DatafileProjectConfig) GetAttributeByKey(key string) (entities.Attribute, error) {
+	if attribute, ok := c.attributeKeyMap[key]; ok {
+		return attribute, nil
+	}
+
+	return entities.Attribute{}, fmt.Errorf(`attribute with key "%s" not found`, key)
+}
+
+// GetAttributeKeyByID returns the attribute key for the given ID
+func (c DatafileProjectConfig) GetAttributeKeyByID(id string) (string, error) {
+	if key, ok := c.attributeIDToKeyMap[id]; ok {
+		return key, nil
+	}
+
+	return "", fmt.Errorf(`attribute with ID "%s" not found`, id)
+}
+
 // GetBotFiltering returns botFiltering
 func (c DatafileProjectConfig) GetBotFiltering() bool {
 	return c.botFiltering
@@ -163,17 +183,6 @@ func (c DatafileProjectConfig) GetVariableByKey(featureKey, variableKey string)
 	return variable, err
 }
 
-// GetAttributeByKey returns the attribute with the given key
-func (c DatafileProjectConfig) GetAttributeByKey(key string) (entities.Attribute, error) {
-	if attributeID, ok := c.attributeKeyToIDMap[key]; ok {
-		if attribute, ok := c.attributeMap[attributeID]; ok {
-			return attribute, nil
-		}
-	}
-
-	return entities.Attribute{}, fmt.Errorf(`attribute with key "%s" not found`, key)
-}
-
 // GetFeatureList returns an array of all the features
 func (c DatafileProjectConfig) GetFeatureList() (featureList []entities.Feature) {
 	for _, feature := range c.featureMap {
@@ -300,6 +309,14 @@ func NewDatafileProjectConfig(jsonDatafile []byte, logger logging.OptimizelyLogP
 	audienceMap, audienceSegmentList := mappers.MapAudiences(append(datafile.TypedAudiences, datafile.Audiences...))
 	flagVariationsMap := mappers.MapFlagVariations(featureMap)
 
+	attributeKeyMap := make(map[string]entities.Attribute)
+	attributeIDToKeyMap := make(map[string]string)
+
+	for id, attribute := range attributeMap {
+		attributeIDToKeyMap[id] = attribute.Key
+		attributeKeyMap[attribute.Key] = attribute
+	}
+
 	config := &DatafileProjectConfig{
 		hostForODP:           hostForODP,
 		publicKeyForODP:      publicKeyForODP,
@@ -325,6 +342,8 @@ func NewDatafileProjectConfig(jsonDatafile []byte, logger logging.OptimizelyLogP
 		rolloutMap:           rolloutMap,
 		sendFlagDecisions:    datafile.SendFlagDecisions,
 		flagVariationsMap:    flagVariationsMap,
+		attributeKeyMap:      attributeKeyMap,
+		attributeIDToKeyMap:  attributeIDToKeyMap,
 	}
 
 	logger.Info("Datafile is valid.")

pkg/config/datafileprojectconfig/config_test.go

@@ -316,18 +316,24 @@ func TestGetVariableByKeyWithMissingVariableError(t *testing.T) {
 }
 
 func TestGetAttributeByKey(t *testing.T) {
-	id := "id"
 	key := "key"
-	attributeKeyToIDMap := make(map[string]string)
-	attributeKeyToIDMap[key] = id
-
 	attribute := entities.Attribute{
 		Key: key,
 	}
+
+	// The old and new mappings to ensure backward compatibility
+	attributeKeyMap := make(map[string]entities.Attribute)
+	attributeKeyMap[key] = attribute
+
+	id := "id"
+	attributeKeyToIDMap := make(map[string]string)
+	attributeKeyToIDMap[key] = id
+
 	attributeMap := make(map[string]entities.Attribute)
 	attributeMap[id] = attribute
 
 	config := &DatafileProjectConfig{
+		attributeKeyMap:     attributeKeyMap,
 		attributeKeyToIDMap: attributeKeyToIDMap,
 		attributeMap:        attributeMap,
 	}
@@ -568,3 +574,129 @@ func TestGetFlagVariationsMap(t *testing.T) {
 	assert.NotNil(t, flagVariationsMap["feature_3"])
 	assert.Len(t, flagVariationsMap["feature_3"], 0)
 }
+
+func TestCmabExperiments(t *testing.T) {
+	// Load the decide-test-datafile.json
+	absPath, _ := filepath.Abs("../../../test-data/decide-test-datafile.json")
+	datafile, err := os.ReadFile(absPath)
+	assert.NoError(t, err)
+
+	// Parse the datafile to modify it
+	var datafileJSON map[string]interface{}
+	err = json.Unmarshal(datafile, &datafileJSON)
+	assert.NoError(t, err)
+
+	// Add CMAB to the first experiment with traffic allocation as an integer
+	experiments := datafileJSON["experiments"].([]interface{})
+	exp0 := experiments[0].(map[string]interface{})
+	exp0["cmab"] = map[string]interface{}{
+		"attributes":        []string{"808797688", "808797689"},
+		"trafficAllocation": 5000, // Changed from array to integer
+	}
+
+	// Convert back to JSON
+	modifiedDatafile, err := json.Marshal(datafileJSON)
+	assert.NoError(t, err)
+
+	// Create project config from modified datafile
+	config, err := NewDatafileProjectConfig(modifiedDatafile, logging.GetLogger("", "DatafileProjectConfig"))
+	assert.NoError(t, err)
+
+	// Get the experiment key from the datafile
+	exp0Key := exp0["key"].(string)
+
+	// Test that Cmab fields are correctly mapped for experiment 0
+	experiment0, err := config.GetExperimentByKey(exp0Key)
+	assert.NoError(t, err)
+	assert.NotNil(t, experiment0.Cmab)
+	if experiment0.Cmab != nil {
+		// Test attribute IDs
+		assert.Equal(t, 2, len(experiment0.Cmab.AttributeIds))
+		assert.Contains(t, experiment0.Cmab.AttributeIds, "808797688")
+		assert.Contains(t, experiment0.Cmab.AttributeIds, "808797689")
+
+		// Test traffic allocation as integer
+		assert.Equal(t, 5000, experiment0.Cmab.TrafficAllocation)
+	}
+}
+
+func TestCmabExperimentsNil(t *testing.T) {
+	// Load the decide-test-datafile.json (which doesn't have CMAB by default)
+	absPath, _ := filepath.Abs("../../../test-data/decide-test-datafile.json")
+	datafile, err := os.ReadFile(absPath)
+	assert.NoError(t, err)
+
+	// Create project config from the original datafile
+	config, err := NewDatafileProjectConfig(datafile, logging.GetLogger("", "DatafileProjectConfig"))
+	assert.NoError(t, err)
+
+	// Parse the datafile to get experiment keys
+	var datafileJSON map[string]interface{}
+	err = json.Unmarshal(datafile, &datafileJSON)
+	assert.NoError(t, err)
+
+	experiments := datafileJSON["experiments"].([]interface{})
+	exp0 := experiments[0].(map[string]interface{})
+	exp0Key := exp0["key"].(string)
+
+	// Test that Cmab field is nil for experiment 0
+	experiment0, err := config.GetExperimentByKey(exp0Key)
+	assert.NoError(t, err)
+	assert.Nil(t, experiment0.Cmab, "CMAB field should be nil when not present in datafile")
+
+	// Test another experiment if available
+	if len(experiments) > 1 {
+		exp1 := experiments[1].(map[string]interface{})
+		exp1Key := exp1["key"].(string)
+
+		experiment1, err := config.GetExperimentByKey(exp1Key)
+		assert.NoError(t, err)
+		assert.Nil(t, experiment1.Cmab, "CMAB field should be nil when not present in datafile")
+	}
+}
+
+func TestGetAttributeKeyByID(t *testing.T) {
+	// Setup
+	id := "id"
+	key := "key"
+	attributeIDToKeyMap := make(map[string]string)
+	attributeIDToKeyMap[id] = key
+
+	config := &DatafileProjectConfig{
+		attributeIDToKeyMap: attributeIDToKeyMap,
+	}
+
+	// Test successful case
+	actual, err := config.GetAttributeKeyByID(id)
+	assert.Nil(t, err)
+	assert.Equal(t, key, actual)
+}
+
+func TestGetAttributeKeyByIDWithMissingIDError(t *testing.T) {
+	// Setup
+	config := &DatafileProjectConfig{}
+
+	// Test error case
+	_, err := config.GetAttributeKeyByID("id")
+	if assert.Error(t, err) {
+		assert.Equal(t, fmt.Errorf(`attribute with ID "id" not found`), err)
+	}
+}
+
+func TestGetAttributeByKeyWithDirectMapping(t *testing.T) {
+	key := "key"
+	attribute := entities.Attribute{
+		Key: key,
+	}
+
+	attributeKeyMap := make(map[string]entities.Attribute)
+	attributeKeyMap[key] = attribute
+
+	config := &DatafileProjectConfig{
+		attributeKeyMap: attributeKeyMap,
+	}
+
+	actual, err := config.GetAttributeByKey(key)
+	assert.Nil(t, err)
+	assert.Equal(t, attribute, actual)
+}

pkg/config/datafileprojectconfig/entities/entities.go

@@ -1,5 +1,5 @@
 /****************************************************************************
- * Copyright 2019,2021-2022, Optimizely, Inc. and contributors              *
+ * Copyright 2019,2021-2025, Optimizely, Inc. and contributors              *
  *                                                                          *
  * Licensed under the Apache License, Version 2.0 (the "License");          *
  * you may not use this file except in compliance with the License.         *
@@ -32,6 +32,14 @@ type Attribute struct {
 	Key string `json:"key"`
 }
 
+// Cmab represents the Contextual Multi-Armed Bandit configuration for an experiment.
+// It contains a list of attribute IDs that are used for the CMAB algorithm and
+// traffic allocation settings for the CMAB implementation.
+type Cmab struct {
+	AttributeIds      []string `json:"attributes"`
+	TrafficAllocation int      `json:"trafficAllocation"`
+}
+
 // Experiment represents an Experiment object from the Optimizely datafile
 type Experiment struct {
 	ID                 string              `json:"id"`
@@ -43,6 +51,7 @@ type Experiment struct {
 	AudienceIds        []string            `json:"audienceIds"`
 	ForcedVariations   map[string]string   `json:"forcedVariations"`
 	AudienceConditions interface{}         `json:"audienceConditions"`
+	Cmab               *Cmab               `json:"cmab,omitempty"` // is optional
 }
 
 // Group represents an Group object from the Optimizely datafile

pkg/config/datafileprojectconfig/mappers/experiment.go

@@ -1,5 +1,5 @@
 /****************************************************************************
- * Copyright 2019,2021, Optimizely, Inc. and contributors                   *
+ * Copyright 2019,2021-2025, Optimizely, Inc. and contributors                   *
  *                                                                          *
  * Licensed under the Apache License, Version 2.0 (the "License");          *
  * you may not use this file except in compliance with the License.         *
@@ -91,6 +91,7 @@ func mapExperiment(rawExperiment datafileEntities.Experiment) entities.Experimen
 		AudienceConditionTree: audienceConditionTree,
 		Whitelist:             rawExperiment.ForcedVariations,
 		IsFeatureExperiment:   false,
+		Cmab:                  mapCmab(rawExperiment.Cmab),
 	}
 
 	for _, variation := range rawExperiment.Variations {
@@ -113,3 +114,15 @@ func MergeExperiments(rawExperiments []datafileEntities.Experiment, rawGroups []
 	}
 	return mergedExperiments
 }
+
+func mapCmab(rawCmab *datafileEntities.Cmab) *entities.Cmab {
+	// handle nil case because cmab is optional and can be nil
+	if rawCmab == nil {
+		return nil
+	}
+
+	return &entities.Cmab{
+		AttributeIds:      rawCmab.AttributeIds,
+		TrafficAllocation: rawCmab.TrafficAllocation,
+	}
+}