1 files changed, 60 insertions, 31 deletions

diff --git a/tools/aconfig/aconfig_storage_write_api/tests/storage_write_api_test.rs b/tools/aconfig/aconfig_storage_write_api/tests/storage_write_api_test.rs
index f6c1bbcd7b..367569def4 100644
--- a/tools/aconfig/aconfig_storage_write_api/tests/storage_write_api_test.rs
+++ b/tools/aconfig/aconfig_storage_write_api/tests/storage_write_api_test.rs
@@ -1,38 +1,17 @@
 #[cfg(not(feature = "cargo"))]
 mod aconfig_storage_write_api_test {
-    use aconfig_storage_file::protos::ProtoStorageFiles;
+    use aconfig_storage_file::{FlagInfoBit, FlagValueType};
+    use aconfig_storage_read_api::flag_info_query::find_flag_attribute;
     use aconfig_storage_read_api::flag_value_query::find_boolean_flag_value;
-    use aconfig_storage_write_api::{mapped_file::get_mapped_file, set_boolean_flag_value};
+    use aconfig_storage_write_api::{
+        map_mutable_storage_file, set_boolean_flag_value, set_flag_has_local_override,
+        set_flag_has_server_override,
+    };
 
-    use protobuf::Message;
     use std::fs::{self, File};
-    use std::io::{Read, Write};
+    use std::io::Read;
     use tempfile::NamedTempFile;
 
-    /// Write storage location record pb to a temp file
-    fn write_storage_record_file(flag_val: &str) -> NamedTempFile {
-        let text_proto = format!(
-            r#"
-files {{
-    version: 0
-    container: "system"
-    package_map: "some_package_map"
-    flag_map: "some_flag_map"
-    flag_val: "{}"
-    timestamp: 12345
-}}
-"#,
-            flag_val
-        );
-        let storage_files: ProtoStorageFiles =
-            protobuf::text_format::parse_from_str(&text_proto).unwrap();
-        let mut binary_proto_bytes = Vec::new();
-        storage_files.write_to_vec(&mut binary_proto_bytes).unwrap();
-        let mut file = NamedTempFile::new().unwrap();
-        file.write_all(&binary_proto_bytes).unwrap();
-        file
-    }
-
     /// Create temp file copy
     fn copy_to_temp_rw_file(source_file: &str) -> NamedTempFile {
         let file = NamedTempFile::new().unwrap();
@@ -48,18 +27,24 @@ files {{
         find_boolean_flag_value(&bytes, offset).unwrap()
     }
 
+    /// Get flag attribute at offset
+    fn get_flag_attribute_at_offset(file: &str, value_type: FlagValueType, offset: u32) -> u8 {
+        let mut f = File::open(file).unwrap();
+        let mut bytes = Vec::new();
+        f.read_to_end(&mut bytes).unwrap();
+        find_flag_attribute(&bytes, value_type, offset).unwrap()
+    }
+
     #[test]
     /// Test to lock down flag value update api
     fn test_boolean_flag_value_update() {
         let flag_value_file = copy_to_temp_rw_file("./flag.val");
         let flag_value_path = flag_value_file.path().display().to_string();
-        let record_pb_file = write_storage_record_file(&flag_value_path);
-        let record_pb_path = record_pb_file.path().display().to_string();
 
         // SAFETY:
         // The safety here is ensured as only this single threaded test process will
         // write to this file
-        let mut file = unsafe { get_mapped_file(&record_pb_path, "system").unwrap() };
+        let mut file = unsafe { map_mutable_storage_file(&flag_value_path).unwrap() };
         for i in 0..8 {
             set_boolean_flag_value(&mut file, i, true).unwrap();
             let value = get_boolean_flag_value_at_offset(&flag_value_path, i);
@@ -70,4 +55,48 @@ files {{
             assert!(!value);
         }
     }
+
+    #[test]
+    /// Test to lock down flag has server override update api
+    fn test_set_flag_has_server_override() {
+        let flag_info_file = copy_to_temp_rw_file("./flag.info");
+        let flag_info_path = flag_info_file.path().display().to_string();
+
+        // SAFETY:
+        // The safety here is ensured as only this single threaded test process will
+        // write to this file
+        let mut file = unsafe { map_mutable_storage_file(&flag_info_path).unwrap() };
+        for i in 0..8 {
+            set_flag_has_server_override(&mut file, FlagValueType::Boolean, i, true).unwrap();
+            let attribute =
+                get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+            assert!((attribute & (FlagInfoBit::HasServerOverride as u8)) != 0);
+            set_flag_has_server_override(&mut file, FlagValueType::Boolean, i, false).unwrap();
+            let attribute =
+                get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+            assert!((attribute & (FlagInfoBit::HasServerOverride as u8)) == 0);
+        }
+    }
+
+    #[test]
+    /// Test to lock down flag has local override update api
+    fn test_set_flag_has_local_override() {
+        let flag_info_file = copy_to_temp_rw_file("./flag.info");
+        let flag_info_path = flag_info_file.path().display().to_string();
+
+        // SAFETY:
+        // The safety here is ensured as only this single threaded test process will
+        // write to this file
+        let mut file = unsafe { map_mutable_storage_file(&flag_info_path).unwrap() };
+        for i in 0..8 {
+            set_flag_has_local_override(&mut file, FlagValueType::Boolean, i, true).unwrap();
+            let attribute =
+                get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+            assert!((attribute & (FlagInfoBit::HasLocalOverride as u8)) != 0);
+            set_flag_has_local_override(&mut file, FlagValueType::Boolean, i, false).unwrap();
+            let attribute =
+                get_flag_attribute_at_offset(&flag_info_path, FlagValueType::Boolean, i);
+            assert!((attribute & (FlagInfoBit::HasLocalOverride as u8)) == 0);
+        }
+    }
 }