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

* chore: Bump engine version

* feat: Identity overrides in local evaluation mode

* fix: Update environment on client init

Author: khvn26

Cargo.toml

@@ -18,12 +18,12 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 reqwest = { version = "0.11", features = ["json", "blocking"] }
 url = "2.1"
-chrono = { version = "0.4"}
+chrono = { version = "0.4" }
 log = "0.4"
 flume = "0.10.14"
 
-flagsmith-flag-engine = "0.3.0"
+flagsmith-flag-engine = "0.4.0"
 
 [dev-dependencies]
 httpmock = "0.6"
-rstest = "0.12.0"
+rstest = "0.12.0"

src/flagsmith/mod.rs

@@ -10,6 +10,7 @@ use flagsmith_flag_engine::segments::Segment;
 use log::debug;
 use reqwest::header::{self, HeaderMap};
 use serde_json::json;
+use std::collections::HashMap;
 use std::sync::mpsc::{self, SyncSender, TryRecvError};
 use std::sync::{Arc, Mutex};
 use std::{thread, time::Duration};
@@ -62,6 +63,7 @@ pub struct Flagsmith {
 
 struct DataStore {
     environment: Option<Environment>,
+    identities_with_overrides_by_identifier: HashMap<String, Identity>,
 }
 
 impl Flagsmith {
@@ -109,7 +111,7 @@ impl Flagsmith {
 
         // Put the environment model behind mutex to
         // to share it safely between threads
-        let ds = Arc::new(Mutex::new(DataStore { environment: None }));
+        let ds = Arc::new(Mutex::new(DataStore { environment: None, identities_with_overrides_by_identifier: HashMap::new() }));
         let (tx, rx) = mpsc::sync_channel::<u32>(1);
 
         let flagsmith = Flagsmith {
@@ -141,6 +143,10 @@ impl Flagsmith {
             flagsmith.options.environment_refresh_interval_mills;
 
         if flagsmith.options.enable_local_evaluation {
+            // Update environment once...
+            update_environment(&client, &ds, &environment_url).unwrap();
+            
+            // ...and continue updating in the background
             let ds = Arc::clone(&ds);
             thread::spawn(move || loop {
                 match rx.try_recv() {
@@ -150,14 +156,8 @@ impl Flagsmith {
                     }
                     Err(TryRecvError::Empty) => {}
                 }
-
-                let environment = Some(
-                    get_environment_from_api(&client, environment_url.clone())
-                        .expect("updating environment document failed"),
-                );
-                let mut data = ds.lock().unwrap();
-                data.environment = environment;
                 thread::sleep(Duration::from_millis(environment_refresh_interval_mills));
+                update_environment(&client, &ds, &environment_url).unwrap();
             });
         }
         return flagsmith;
@@ -201,7 +201,7 @@ impl Flagsmith {
         let traits = traits.unwrap_or(vec![]);
         if data.environment.is_some() {
             let environment = data.environment.as_ref().unwrap();
-            return self.get_identity_flags_from_document(environment, identifier, traits);
+            return self.get_identity_flags_from_document(environment, &data.identities_with_overrides_by_identifier, identifier, traits);
         }
         return self.default_handler_if_err(self.get_identity_flags_from_api(identifier, traits));
     }
@@ -219,8 +219,9 @@ impl Flagsmith {
             ));
         }
         let environment = data.environment.as_ref().unwrap();
+        let identities_with_overrides_by_identifier = &data.identities_with_overrides_by_identifier;
         let identity_model =
-            self.build_identity_model(environment, identifier, traits.clone().unwrap_or(vec![]))?;
+            self.get_identity_model(&environment, &identities_with_overrides_by_identifier, identifier, traits.clone().unwrap_or(vec![]))?;
         let segments = get_identity_segments(environment, &identity_model, traits.as_ref());
         return Ok(segments);
     }
@@ -254,21 +255,17 @@ impl Flagsmith {
         );
     }
     pub fn update_environment(&mut self) -> Result<(), error::Error> {
-        let mut data = self.datastore.lock().unwrap();
-        data.environment = Some(get_environment_from_api(
-            &self.client,
-            self.environment_url.clone(),
-        )?);
-        return Ok(());
+        return update_environment(&self.client, &self.datastore, &self.environment_url);
     }
 
     fn get_identity_flags_from_document(
         &self,
         environment: &Environment,
+        identities_with_overrides_by_identifier: &HashMap<String, Identity>,
         identifier: &str,
         traits: Vec<Trait>,
     ) -> Result<Flags, error::Error> {
-        let identity = self.build_identity_model(environment, identifier, traits.clone())?;
+        let identity = self.get_identity_model(environment, identities_with_overrides_by_identifier, identifier, traits.clone())?;
         let feature_states =
             engine::get_identity_feature_states(environment, &identity, Some(traits.as_ref()));
         let flags = Flags::from_feature_states(
@@ -280,15 +277,23 @@ impl Flagsmith {
         return Ok(flags);
     }
 
-    fn build_identity_model(
+    fn get_identity_model(
         &self,
         environment: &Environment,
+        identities_with_overrides_by_identifier: &HashMap<String, Identity>,
         identifier: &str,
         traits: Vec<Trait>,
     ) -> Result<Identity, error::Error> {
-        let mut identity = Identity::new(identifier.to_string(), environment.api_key.clone());
+        let mut identity: Identity;
+
+        if identities_with_overrides_by_identifier.contains_key(identifier) {
+            identity = identities_with_overrides_by_identifier.get(identifier).unwrap().clone();
+        } else {
+            identity = Identity::new(identifier.to_string(), environment.api_key.clone());
+        }
+
         identity.identity_traits = traits;
-        Ok(identity)
+        return Ok(identity.to_owned())
     }
     fn get_identity_flags_from_api(
         &self,
@@ -358,6 +363,23 @@ fn get_environment_from_api(
     return Ok(environment);
 }
 
+fn update_environment(
+    client: &reqwest::blocking::Client,
+    datastore: &Arc<Mutex<DataStore>>,
+    environment_url: &String,
+) -> Result<(), error::Error> {
+    let mut data = datastore.lock().unwrap();
+    let environment = Some(get_environment_from_api(
+        &client,
+        environment_url.clone(),
+    )?);
+    for identity in &environment.as_ref().unwrap().identity_overrides {
+        data.identities_with_overrides_by_identifier.insert(identity.identifier.clone(), identity.clone());
+    }
+    data.environment = environment;
+    return Ok(());
+}
+
 fn get_json_response(
     client: &reqwest::blocking::Client,
     method: reqwest::Method,
@@ -385,24 +407,87 @@ mod tests {
     use httpmock::prelude::*;
 
     static ENVIRONMENT_JSON: &str = r#"{
-            "api_key": "B62qaMZNwfiqT76p38ggrQ",
-            "project": {
-                "name": "Test project",
-                "organisation": {
-                    "feature_analytics": false,
-                    "name": "Test Org",
-                    "id": 1,
-                    "persist_trait_data": true,
-                    "stop_serving_flags": false
+        "api_key": "B62qaMZNwfiqT76p38ggrQ",
+        "project": {
+          "name": "Test project",
+          "organisation": {
+            "feature_analytics": false,
+            "name": "Test Org",
+            "id": 1,
+            "persist_trait_data": true,
+            "stop_serving_flags": false
+          },
+          "segments": [],
+          "id": 1,
+          "hide_disabled_flags": false
+        },
+        "segment_overrides": [],
+        "id": 1,
+        "feature_states": [
+          {
+            "multivariate_feature_state_values": [],
+            "feature_state_value": "some-value",
+            "id": 1,
+            "featurestate_uuid": "40eb539d-3713-4720-bbd4-829dbef10d51",
+            "feature": {
+              "name": "some_feature",
+              "type": "STANDARD",
+              "id": 1
+            },
+            "segment_id": null,
+            "enabled": true
+          },
+          {
+            "feature": {
+              "id": 83755,
+              "name": "test_mv",
+              "type": "MULTIVARIATE"
+            },
+            "enabled": false,
+            "django_id": 482285,
+            "feature_segment": null,
+            "featurestate_uuid": "c3af5fbf-39ba-422c-a846-f2fea952b37c",
+            "feature_state_value": "1111",
+            "multivariate_feature_state_values": [
+              {
+                "multivariate_feature_option": {
+                  "value": "8888",
+                  "id": 11516
                 },
+                "percentage_allocation": 100.0,
+                "id": 38451,
+                "mv_fs_value_uuid": "a4299c73-2430-47e4-9185-42accd01686c"
+              }
+            ]
+          }
+        ],
+        "updated_at": "2023-07-14 16:12:00.000000",
+        "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,
-                "hide_disabled_flags": false,
-                "segments": []
-            },
-            "segment_overrides": [],
-            "id": 1,
-            "feature_states": []
-    }"#;
+                "feature": {
+                  "id": 1,
+                  "name": "some_feature",
+                  "type": "STANDARD"
+                },
+                "featurestate_uuid": "1bddb9a5-7e59-42c6-9be9-625fa369749f",
+                "feature_state_value": "some-overridden-value",
+                "enabled": false,
+                "environment": 1,
+                "identity": null,
+                "feature_segment": null
+              }
+            ]
+          }
+        ]
+      }"#;
 
     #[test]
     fn client_implements_send_and_sync() {
@@ -472,4 +557,37 @@ mod tests {
         // for each subsequent refresh
         api_mock.assert_hits(3);
     }
+
+    #[test]
+    fn test_local_evaluation_identity_override_evaluate_expected() {
+        // Given
+        let environment_key = "ser.test_environment_key";
+        let response_body: serde_json::Value = serde_json::from_str(ENVIRONMENT_JSON).unwrap();
+
+        let mock_server = MockServer::start();
+        mock_server.mock(|when, then| {
+            when.method(GET)
+                .path("/api/v1/environment-document/")
+                .header("X-Environment-Key", environment_key);
+            then.status(200).json_body(response_body);
+        });
+
+        let url = mock_server.url("/api/v1/");
+
+        let flagsmith_options = FlagsmithOptions {
+            api_url: url,
+            environment_refresh_interval_mills: 100,
+            enable_local_evaluation: true,
+            ..Default::default()
+        };
+
+        // When
+        let mut _flagsmith = Flagsmith::new(environment_key.to_string(), flagsmith_options);
+
+        // Then
+        let flags = _flagsmith.get_environment_flags();
+        let identity_flags = _flagsmith.get_identity_flags("overridden-id", None);
+        assert_eq!(flags.unwrap().get_feature_value_as_string("some_feature").unwrap().to_owned(), "some-value");
+        assert_eq!(identity_flags.unwrap().get_feature_value_as_string("some_feature").unwrap().to_owned(), "some-overridden-value");
+    }
 }