feat: Identity overrides in local evaluation mode (#121)

* feat: Identity overrides in local evaluation mode
Flagsmith · Apr 5, 2024 · e592264 · e592264
1 parent 569276f
commit e592264
Show file tree

Hide file tree

Showing 5 changed files with 104 additions and 16 deletions.
diff --git a/client.go b/client.go
@@ -22,7 +22,8 @@ type Client struct {
 	apiKey string
 	config config
 
-	environment atomic.Value
+	environment             atomic.Value
+	identitiesWithOverrides atomic.Value
 
 	analyticsProcessor *AnalyticsProcessor
 	defaultFlagHandler func(string) (Flag, error)
@@ -138,7 +139,7 @@ func (c *Client) GetIdentityFlags(ctx context.Context, identifier string, traits
 // Returns an array of segments that the given identity is part of.
 func (c *Client) GetIdentitySegments(identifier string, traits []*Trait) ([]*segments.SegmentModel, error) {
 	if env, ok := c.environment.Load().(*environments.EnvironmentModel); ok {
-		identity := buildIdentityModel(identifier, env.APIKey, traits)
+		identity := c.getIdentityModel(identifier, env.APIKey, traits)
 		return flagengine.GetIdentitySegments(env, &identity), nil
 	}
 	return nil, &FlagsmithClientError{msg: "flagsmith: Local evaluation required to obtain identity segments"}
@@ -214,7 +215,7 @@ func (c *Client) getIdentityFlagsFromEnvironment(identifier string, traits []*Tr
 	if !ok {
 		return Flags{}, fmt.Errorf("flagsmith: local environment has not yet been updated")
 	}
-	identity := buildIdentityModel(identifier, env.APIKey, traits)
+	identity := c.getIdentityModel(identifier, env.APIKey, traits)
 	featureStates := flagengine.GetIdentityFeatureStates(env, &identity)
 	flags := makeFlagsFromFeatureStates(
 		featureStates,
@@ -276,15 +277,28 @@ func (c *Client) UpdateEnvironment(ctx context.Context) error {
 		return errors.New(e["detail"])
 	}
 	c.environment.Store(&env)
+	identitiesWithOverrides := make(map[string]identities.IdentityModel)
+	for _, id := range env.IdentityOverrides {
+		identitiesWithOverrides[id.Identifier] = *id
+	}
+	c.identitiesWithOverrides.Store(identitiesWithOverrides)
 
 	return nil
 }
 
-func buildIdentityModel(identifier string, apiKey string, traits []*Trait) identities.IdentityModel {
+func (c *Client) getIdentityModel(identifier string, apiKey string, traits []*Trait) identities.IdentityModel {
 	identityTraits := make([]*TraitModel, len(traits))
 	for i, trait := range traits {
 		identityTraits[i] = trait.ToTraitModel()
 	}
+
+	identitiesWithOverrides, _ := c.identitiesWithOverrides.Load().(map[string]identities.IdentityModel)
+	identity, ok := identitiesWithOverrides[identifier]
+	if ok {
+		identity.IdentityTraits = identityTraits
+		return identity
+	}
+
 	return identities.IdentityModel{
 		Identifier:        identifier,
 		IdentityTraits:    identityTraits,

diff --git a/client_test.go b/client_test.go
@@ -244,6 +244,32 @@ func TestGetIdentityFlagsUseslocalEnvironmentWhenAvailable(t *testing.T) {
 	assert.Equal(t, fixtures.Feature1Value, allFlags[0].Value)
 }
 
+func TestGetIdentityFlagsUseslocalOverridesWhenAvailable(t *testing.T) {
+	// Given
+	ctx := context.Background()
+	server := httptest.NewServer(http.HandlerFunc(fixtures.EnvironmentDocumentHandler))
+	defer server.Close()
+	// When
+	client := flagsmith.NewClient(fixtures.EnvironmentAPIKey, flagsmith.WithLocalEvaluation(ctx),
+		flagsmith.WithBaseURL(server.URL+"/api/v1/"))
+	err := client.UpdateEnvironment(ctx)
+
+	// Then
+	assert.NoError(t, err)
+
+	flags, err := client.GetIdentityFlags(ctx, "overridden-id", nil)
+
+	assert.NoError(t, err)
+
+	allFlags := flags.AllFlags()
+
+	assert.Equal(t, 1, len(allFlags))
+
+	assert.Equal(t, fixtures.Feature1Name, allFlags[0].FeatureName)
+	assert.Equal(t, fixtures.Feature1ID, allFlags[0].FeatureID)
+	assert.Equal(t, fixtures.Feature1OverriddenValue, allFlags[0].Value)
+}
+
 func TestGetIdentityFlagsCallsAPIWhenLocalEnvironmentNotAvailableWithTraits(t *testing.T) {
 	// Given
 	ctx := context.Background()

diff --git a/fixtures/environment.json b/fixtures/environment.json
@@ -54,5 +54,30 @@
             "segment_id": null,
             "enabled": true
         }
+    ],
+    "identity_overrides": [
+        {
+            "identifier": "overridden-id",
+            "identity_uuid": "0f21cde8-63c5-4e50-baca-87897fa6cd01",
+            "created_date": "2019-08-27T14:53:45.698555Z",
+            "updated_at": "2023-07-14 16:12:00.000000",
+            "environment_api_key": "B62qaMZNwfiqT76p38ggrQ",
+            "identity_features": [
+                {
+                    "id": 1,
+                    "feature": {
+                        "id": 1,
+                        "name": "feature_1",
+                        "type": "STANDARD"
+                    },
+                    "featurestate_uuid": "1bddb9a5-7e59-42c6-9be9-625fa369749f",
+                    "feature_state_value": "some-overridden-value",
+                    "enabled": false,
+                    "environment": 1,
+                    "identity": null,
+                    "feature_segment": null
+                }
+            ]
+        }
     ]
-}
+}
diff --git a/fixtures/fixture.go b/fixtures/fixture.go
@@ -11,6 +11,8 @@ const Feature1Value = "some_value"
 const Feature1Name = "feature_1"
 const Feature1ID = 1
 
+const Feature1OverriddenValue = "some-overridden-value"
+
 const EnvironmentJson = `
 {
 	"api_key": "B62qaMZNwfiqT76p38ggrQ",
@@ -58,7 +60,32 @@ const EnvironmentJson = `
 		},
 		"segment_id": null,
 		"enabled": true
-	}]
+	}],
+    "identity_overrides": [
+        {
+            "identifier": "overridden-id",
+            "identity_uuid": "0f21cde8-63c5-4e50-baca-87897fa6cd01",
+            "created_date": "2019-08-27T14:53:45.698555Z",
+            "updated_at": "2023-07-14 16:12:00.000000",
+            "environment_api_key": "B62qaMZNwfiqT76p38ggrQ",
+            "identity_features": [
+                {
+                    "id": 1,
+                    "feature": {
+                        "id": 1,
+                        "name": "feature_1",
+                        "type": "STANDARD"
+                    },
+                    "featurestate_uuid": "1bddb9a5-7e59-42c6-9be9-625fa369749f",
+                    "feature_state_value": "some-overridden-value",
+                    "enabled": false,
+                    "environment": 1,
+                    "identity": null,
+                    "feature_segment": null
+                }
+            ]
+        }
+    ]
 }
 `
 

diff --git a/flagengine/environments/models.go b/flagengine/environments/models.go
@@ -1,19 +1,15 @@
 package environments
 
 import (
-	"github.com/Flagsmith/flagsmith-go-client/v3/flagengine/environments/integrations"
 	"github.com/Flagsmith/flagsmith-go-client/v3/flagengine/features"
+	"github.com/Flagsmith/flagsmith-go-client/v3/flagengine/identities"
 	"github.com/Flagsmith/flagsmith-go-client/v3/flagengine/projects"
 )
 
 type EnvironmentModel struct {
-	ID            int                           `json:"id"`
-	APIKey        string                        `json:"api_key"`
-	Project       *projects.ProjectModel        `json:"project"`
-	FeatureStates []*features.FeatureStateModel `json:"feature_states"`
-
-	AmplitudeConfig *integrations.IntegrationModel `json:"amplitude_config"`
-	SegmentConfig   *integrations.IntegrationModel `json:"segment_config"`
-	MixpanelConfig  *integrations.IntegrationModel `json:"mixpanel_config"`
-	HeapConfig      *integrations.IntegrationModel `json:"heap_config"`
+	ID                int                           `json:"id"`
+	APIKey            string                        `json:"api_key"`
+	Project           *projects.ProjectModel        `json:"project"`
+	FeatureStates     []*features.FeatureStateModel `json:"feature_states"`
+	IdentityOverrides []*identities.IdentityModel   `json:"identity_overrides"`
 }