Initialise user_type fields

open-chat-labs · Jul 26, 2024 · f6ff30d · f6ff30d
1 parent 38bfdd6
commit f6ff30d
Show file tree

Hide file tree

Showing 9 changed files with 66 additions and 5 deletions.
diff --git a/backend/canisters/community/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/community/impl/src/lifecycle/post_upgrade.rs
@@ -1,14 +1,15 @@
 use crate::jobs::import_groups::finalize_group_import;
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::get_upgrades_memory;
-use crate::{read_state, Data};
+use crate::{mutate_state, read_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use community_canister::post_upgrade::Args;
 use ic_cdk::post_upgrade;
 use instruction_counts_log::InstructionCountFunctionId;
 use stable_memory::get_reader;
 use tracing::info;
+use utils::consts::OPENCHAT_BOT_USER_ID;
 
 #[post_upgrade]
 #[trace]
@@ -31,6 +32,12 @@ fn post_upgrade(args: Args) {
 
     info!(version = %args.wasm_version, "Post-upgrade complete");
 
+    mutate_state(|state| {
+        let oc_controlled_bot_users = vec![OPENCHAT_BOT_USER_ID, state.data.proposals_bot_user_id];
+        state.data.members.set_user_types(&oc_controlled_bot_users);
+        state.data.channels.set_user_types(&oc_controlled_bot_users);
+    });
+
     read_state(|state| {
         let now = state.env.now();
         state

diff --git a/backend/canisters/community/impl/src/model/channels.rs b/backend/canisters/community/impl/src/model/channels.rs
@@ -28,6 +28,12 @@ pub struct Channel {
 }
 
 impl Channels {
+    pub fn set_user_types(&mut self, oc_controlled_bots: &[UserId]) {
+        for channel in self.channels.values_mut() {
+            channel.chat.set_user_types(oc_controlled_bots);
+        }
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         community_id: CommunityId,

diff --git a/backend/canisters/community/impl/src/model/members.rs b/backend/canisters/community/impl/src/model/members.rs
@@ -23,6 +23,13 @@ pub struct CommunityMembers {
 }
 
 impl CommunityMembers {
+    pub fn set_user_types(&mut self, oc_controlled_bot_users: &[UserId]) {
+        for (user_id, member) in self.members.iter_mut().filter(|(_, m)| m.is_bot) {
+            member.user_type =
+                if oc_controlled_bot_users.contains(user_id) { UserType::OcControlledBot } else { UserType::Bot };
+        }
+    }
+
     pub fn new(
         creator_principal: Principal,
         creator_user_id: UserId,

diff --git a/backend/canisters/group/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/group/impl/src/lifecycle/post_upgrade.rs
@@ -1,13 +1,14 @@
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::get_upgrades_memory;
-use crate::{read_state, Data};
+use crate::{mutate_state, read_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use group_canister::post_upgrade::Args;
 use ic_cdk::post_upgrade;
 use instruction_counts_log::InstructionCountFunctionId;
 use stable_memory::get_reader;
 use tracing::info;
+use utils::consts::OPENCHAT_BOT_USER_ID;
 
 #[post_upgrade]
 #[trace]
@@ -24,6 +25,11 @@ fn post_upgrade(args: Args) {
 
     info!(version = %args.wasm_version, "Post-upgrade complete");
 
+    mutate_state(|state| {
+        let oc_controlled_bot_users = vec![OPENCHAT_BOT_USER_ID, state.data.proposals_bot_user_id];
+        state.data.chat.set_user_types(&oc_controlled_bot_users);
+    });
+
     read_state(|state| {
         let now = state.env.now();
         state

diff --git a/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/user/impl/src/lifecycle/post_upgrade.rs
@@ -1,6 +1,6 @@
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::get_upgrades_memory;
-use crate::{read_state, Data};
+use crate::{mutate_state, read_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
@@ -28,6 +28,12 @@ fn post_upgrade(args: Args) {
 
     info!(version = %args.wasm_version, "Post-upgrade complete");
 
+    mutate_state(|state| {
+        for chat in state.data.direct_chats.iter_mut() {
+            chat.set_user_type();
+        }
+    });
+
     read_state(|state| {
         if state.data.direct_chats.len() <= 1
             && state.data.group_chats.len() == 0

diff --git a/backend/canisters/user/impl/src/model/direct_chat.rs b/backend/canisters/user/impl/src/model/direct_chat.rs
@@ -8,6 +8,7 @@ use types::{
     TimestampMillis, Timestamped, UserId, UserType,
 };
 use user_canister::SendMessageArgs;
+use utils::consts::OPENCHAT_BOT_USER_ID;
 
 #[derive(Serialize, Deserialize)]
 pub struct DirectChat {
@@ -26,6 +27,12 @@ pub struct DirectChat {
 }
 
 impl DirectChat {
+    pub fn set_user_type(&mut self) {
+        if self.is_bot {
+            self.user_type = if self.them == OPENCHAT_BOT_USER_ID { UserType::OcControlledBot } else { UserType::Bot };
+        }
+    }
+
     pub fn new(
         them: UserId,
         user_type: UserType,

diff --git a/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs
@@ -1,11 +1,13 @@
 use crate::lifecycle::{init_env, init_state};
 use crate::memory::get_upgrades_memory;
-use crate::Data;
+use crate::{mutate_state, Data};
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
+use local_user_index_canister::UserRegistered;
 use stable_memory::get_reader;
 use tracing::info;
+use types::UserType;
 use user_index_canister::post_upgrade::Args;
 use utils::cycles::init_cycles_dispenser_client;
 
@@ -23,5 +25,19 @@ fn post_upgrade(args: Args) {
     init_cycles_dispenser_client(data.cycles_dispenser_canister_id, data.test_mode);
     init_state(env, data, args.wasm_version);
 
+    mutate_state(|state| {
+        state.push_event_to_all_local_user_indexes(
+            local_user_index_canister::Event::UserRegistered(UserRegistered {
+                user_id: state.data.proposals_bot_canister_id.into(),
+                user_principal: state.data.proposals_bot_canister_id,
+                username: "ProposalsBot".to_string(),
+                is_bot: true,
+                user_type: UserType::OcControlledBot,
+                referred_by: None,
+            }),
+            None,
+        );
+    });
+
     info!(version = %args.wasm_version, "Post-upgrade complete");
 }
diff --git a/backend/libraries/group_chat_core/src/lib.rs b/backend/libraries/group_chat_core/src/lib.rs
@@ -60,6 +60,12 @@ pub struct GroupChatCore {
 
 #[allow(clippy::too_many_arguments)]
 impl GroupChatCore {
+    pub fn set_user_types(&mut self, oc_controlled_bots: &[UserId]) {
+        for (user_id, member) in self.members.members.iter_mut().filter(|(_, m)| m.is_bot) {
+            member.user_type = if oc_controlled_bots.contains(user_id) { UserType::OcControlledBot } else { UserType::Bot };
+        }
+    }
+
     pub fn new(
         chat: MultiUserChat,
         created_by: UserId,

diff --git a/backend/libraries/group_chat_core/src/members.rs b/backend/libraries/group_chat_core/src/members.rs
@@ -513,7 +513,7 @@ mod tests {
         let member_bytes = msgpack::serialize_then_unwrap(&member);
         let member_bytes_len = member_bytes.len();
 
-        assert_eq!(member_bytes_len, 120);
+        assert_eq!(member_bytes_len, 127);
 
         let _deserialized: GroupMemberInternal = msgpack::deserialize_then_unwrap(&member_bytes);
     }