Fix local_user_index_canister function (#6052)

open-chat-labs · Jul 17, 2024 · f289a88 · f289a88
1 parent a76414c
commit f289a88
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/backend/canisters/user_index/impl/Cargo.toml b/backend/canisters/user_index/impl/Cargo.toml
@@ -57,6 +57,7 @@ sha2 = { workspace = true }
 stable_memory = { path = "../../../libraries/stable_memory" }
 storage_index_canister = { path = "../../storage_index/api" }
 storage_index_canister_c2c_client = { path = "../../storage_index/c2c_client" }
+test-case = { workspace = true }
 testing = { path = "../../../libraries/testing" }
 time = { workspace = true }
 tracing = { workspace = true }

diff --git a/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs b/backend/canisters/user_index/impl/src/lifecycle/post_upgrade.rs
@@ -5,6 +5,7 @@ use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
 use ic_cdk::post_upgrade;
 use stable_memory::get_reader;
+use test_case::test_case;
 use tracing::info;
 use types::CanisterId;
 use user_index_canister::post_upgrade::Args;
@@ -47,17 +48,28 @@ fn post_upgrade(args: Args) {
 fn local_user_index_canister(canister_id: CanisterId, test_mode: bool) -> CanisterId {
     let bytes = canister_id.as_slice();
     if bytes > [0, 0, 0, 0, 2, 32, 0, 0, 1, 1].as_slice() && bytes < [0, 0, 0, 0, 2, 48, 0, 0, 1, 1].as_slice() {
+        return CanisterId::from_text("aboy3-giaaa-aaaar-aaaaq-cai").unwrap();
+    }
+    if bytes > [0, 0, 0, 0, 0, 160, 0, 0, 1, 1].as_slice() && bytes < [0, 0, 0, 0, 0, 176, 0, 0, 1, 1].as_slice() {
         return if test_mode {
             CanisterId::from_text("pecvb-tqaaa-aaaaf-bhdiq-cai").unwrap()
         } else {
             CanisterId::from_text("nq4qv-wqaaa-aaaaf-bhdgq-cai").unwrap()
         };
     }
-    if bytes > [0, 0, 0, 0, 0, 160, 0, 0, 1, 1].as_slice() && bytes < [0, 0, 0, 0, 0, 176, 0, 0, 1, 1].as_slice() {
-        return CanisterId::from_text("aboy3-giaaa-aaaar-aaaaq-cai").unwrap();
-    }
 
     assert!(test_mode);
     // This will only be reached during tests + local development
     CanisterId::from_text("be2us-64aaa-aaaaa-qaabq-cai").unwrap()
 }
+
+#[test_case("foyjw-giaaa-aaaaf-big5a-cai", true, "pecvb-tqaaa-aaaaf-bhdiq-cai")]
+#[test_case("x2jhd-2yaaa-aaaaf-bif6a-cai", true, "pecvb-tqaaa-aaaaf-bhdiq-cai")]
+#[test_case("ueepn-qiaaa-aaaar-ampla-cai", false, "aboy3-giaaa-aaaar-aaaaq-cai")]
+#[test_case("nfjlt-kaaaa-aaaar-axpxq-cai", false, "aboy3-giaaa-aaaar-aaaaq-cai")]
+#[test_case("7msbr-laaaa-aaaaf-ayvuq-cai", false, "nq4qv-wqaaa-aaaaf-bhdgq-cai")]
+#[test_case("2cpmn-myaaa-aaaaf-ayvia-cai", false, "nq4qv-wqaaa-aaaaf-bhdgq-cai")]
+fn local_user_index_canister_tests(canister_id: &str, test_mode: bool, expected: &str) {
+    let result = local_user_index_canister(CanisterId::from_text(canister_id).unwrap(), test_mode);
+    assert_eq!(result.to_text(), expected);
+}