test(drive): improve identity fetch, balance, and public key hash coverage

packages/rs-drive/src/drive/identity/fetch/balance/mod.rs

@@ -306,4 +306,264 @@ mod tests {
             assert_eq!(negative_balance, 0);
         }
     }
+
+    mod fetch_identity_balance_with_transaction {
+        use super::*;
+        use crate::config::DriveConfig;
+        use crate::util::test_helpers::setup::setup_drive;
+
+        #[test]
+        fn should_return_balance_within_transaction() {
+            let drive = setup_drive(Some(DriveConfig {
+                batching_consistency_verification: true,
+                ..Default::default()
+            }));
+            let platform_version = PlatformVersion::latest();
+
+            let transaction = drive.grove.start_transaction();
+            drive
+                .create_initial_state_structure(Some(&transaction), platform_version)
+                .expect("should create root tree");
+
+            let identity = Identity::random_identity(3, Some(42), platform_version)
+                .expect("expected a random identity");
+
+            let expected_balance = identity.balance();
+
+            drive
+                .add_new_identity(
+                    identity.clone(),
+                    false,
+                    &BlockInfo::default(),
+                    true,
+                    Some(&transaction),
+                    platform_version,
+                )
+                .expect("expected to add identity");
+
+            let balance = drive
+                .fetch_identity_balance(
+                    identity.id().to_buffer(),
+                    Some(&transaction),
+                    platform_version,
+                )
+                .expect("should not error")
+                .expect("should have balance");
+
+            assert_eq!(balance, expected_balance);
+
+            let balance_outside = drive
+                .fetch_identity_balance(identity.id().to_buffer(), None, platform_version)
+                .expect("should not error");
+
+            assert!(balance_outside.is_none());
+        }
+    }
+
+    mod fetch_identity_balance_with_costs_applied {
+        use super::*;
+
+        #[test]
+        fn should_return_actual_balance_with_costs_when_applied() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity = Identity::random_identity(3, Some(42), platform_version)
+                .expect("expected a random identity");
+
+            let expected_balance = identity.balance();
+
+            drive
+                .add_new_identity(
+                    identity.clone(),
+                    false,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("expected to add identity");
+
+            let block_info = BlockInfo::default();
+
+            let (balance, fee_result) = drive
+                .fetch_identity_balance_with_costs(
+                    identity.id().to_buffer(),
+                    &block_info,
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should return balance with costs");
+
+            assert_eq!(balance, Some(expected_balance));
+            assert!(fee_result.processing_fee > 0);
+        }
+
+        #[test]
+        fn should_return_none_with_costs_for_non_existent_identity() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let block_info = BlockInfo::default();
+
+            let (balance, fee_result) = drive
+                .fetch_identity_balance_with_costs(
+                    [0u8; 32],
+                    &block_info,
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should return none with costs");
+
+            assert!(balance.is_none());
+            assert!(fee_result.processing_fee > 0);
+        }
+    }
+
+    mod fetch_identity_balance_include_debt_with_costs {
+        use super::*;
+        use crate::fees::op::LowLevelDriveOperation;
+
+        #[test]
+        fn should_return_balance_with_costs_estimated() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity = create_test_identity(&drive, [0; 32], Some(1), None, platform_version)
+                .expect("expected an identity");
+
+            let added_balance = 1000;
+            drive
+                .add_to_identity_balance(
+                    identity.id().to_buffer(),
+                    added_balance,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should add balance");
+
+            let block_info = BlockInfo::default();
+
+            let (balance, fee_result) = drive
+                .fetch_identity_balance_include_debt_with_costs(
+                    identity.id().to_buffer(),
+                    &block_info,
+                    false,
+                    None,
+                    platform_version,
+                )
+                .expect("should return with costs");
+
+            assert!(fee_result.processing_fee > 0);
+            assert!(balance.is_some());
+        }
+
+        #[test]
+        fn should_return_actual_balance_with_costs_when_applied() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity = create_test_identity(&drive, [0; 32], Some(1), None, platform_version)
+                .expect("expected an identity");
+
+            let added_balance: u64 = 2000;
+            drive
+                .add_to_identity_balance(
+                    identity.id().to_buffer(),
+                    added_balance,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should add balance");
+
+            let block_info = BlockInfo::default();
+
+            let (balance, fee_result) = drive
+                .fetch_identity_balance_include_debt_with_costs(
+                    identity.id().to_buffer(),
+                    &block_info,
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should return with costs");
+
+            assert_eq!(balance, Some(added_balance as i64));
+            assert!(fee_result.processing_fee > 0);
+        }
+
+        #[test]
+        fn should_return_negative_balance_with_costs_for_debt() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity = create_test_identity(&drive, [0; 32], Some(1), None, platform_version)
+                .expect("expected an identity");
+
+            let negative_amount: u64 = 500;
+
+            let batch = vec![drive
+                .update_identity_negative_credit_operation(
+                    identity.id().to_buffer(),
+                    negative_amount,
+                    platform_version,
+                )
+                .expect("expected operation")];
+
+            let mut drive_operations: Vec<LowLevelDriveOperation> = vec![];
+            drive
+                .apply_batch_low_level_drive_operations(
+                    None,
+                    None,
+                    batch,
+                    &mut drive_operations,
+                    &platform_version.drive,
+                )
+                .expect("should apply batch");
+
+            let block_info = BlockInfo::default();
+
+            let (balance, fee_result) = drive
+                .fetch_identity_balance_include_debt_with_costs(
+                    identity.id().to_buffer(),
+                    &block_info,
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should return with costs");
+
+            assert_eq!(balance, Some(-(negative_amount as i64)));
+            assert!(fee_result.processing_fee > 0);
+        }
+    }
+
+    mod fetch_identity_negative_balance_estimated {
+        use super::*;
+
+        #[test]
+        fn should_return_zero_in_estimated_mode_for_non_existent_identity() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let mut drive_operations = vec![];
+            let result = drive
+                .fetch_identity_negative_balance_operations(
+                    [0xffu8; 32],
+                    false,
+                    None,
+                    &mut drive_operations,
+                    platform_version,
+                )
+                .expect("should not error in estimated mode");
+
+            assert_eq!(result, Some(0));
+        }
+    }
 }

packages/rs-drive/src/drive/identity/fetch/contract_keys/mod.rs

@@ -60,3 +60,110 @@ impl Drive {
         }
     }
 }
+
+#[cfg(feature = "server")]
+#[cfg(test)]
+mod tests {
+    use crate::util::test_helpers::setup::setup_drive_with_initial_state_structure;
+    use dpp::identity::Purpose;
+    use dpp::version::PlatformVersion;
+
+    mod fetch_identities_contract_keys {
+        use super::*;
+
+        #[test]
+        fn should_return_empty_map_when_no_contract_keys_exist() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity_ids = [[1u8; 32]];
+            let contract_id = [2u8; 32];
+            let purposes = vec![Purpose::ENCRYPTION];
+
+            // When there are no contract keys bound, the query returns an
+            // empty result (the identity subtree exists but has no contract info).
+            let result = drive.fetch_identities_contract_keys(
+                &identity_ids,
+                &contract_id,
+                None,
+                purposes,
+                None,
+                platform_version,
+            );
+
+            let map = result.expect("expected Ok result for non-existent identity");
+            assert!(
+                map.is_empty(),
+                "expected empty map for non-existent identity"
+            );
+        }
+
+        #[test]
+        fn should_return_empty_for_existing_identity_without_contract_keys() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            use dpp::block::block_info::BlockInfo;
+            use dpp::identity::accessors::IdentityGettersV0;
+            use dpp::identity::Identity;
+
+            let identity = Identity::random_identity(3, Some(42), platform_version)
+                .expect("expected a random identity");
+
+            drive
+                .add_new_identity(
+                    identity.clone(),
+                    false,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("expected to add identity");
+
+            let identity_ids = [identity.id().to_buffer()];
+            let contract_id = [0xabu8; 32];
+            let purposes = vec![Purpose::ENCRYPTION];
+
+            // The identity exists but has no contract-bound keys, so the
+            // query should return an empty result or skip that identity.
+            let result = drive.fetch_identities_contract_keys(
+                &identity_ids,
+                &contract_id,
+                None,
+                purposes,
+                None,
+                platform_version,
+            );
+
+            let map = result.expect("expected Ok result for identity without contract keys");
+            assert!(
+                map.is_empty(),
+                "expected empty map when no contract keys exist"
+            );
+        }
+
+        #[test]
+        fn should_return_empty_for_empty_identity_ids() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identity_ids: [[u8; 32]; 0] = [];
+            let contract_id = [3u8; 32];
+            let purposes = vec![Purpose::ENCRYPTION];
+
+            let result = drive
+                .fetch_identities_contract_keys(
+                    &identity_ids,
+                    &contract_id,
+                    None,
+                    purposes,
+                    None,
+                    platform_version,
+                )
+                .expect("should not error for empty ids");
+
+            assert!(result.is_empty());
+        }
+    }
+}