test(drive): address review comments on identity test coverage

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

@@ -188,18 +188,86 @@ mod tests {
         }
     }
 
+    mod fetch_identity_balance_include_debt_combined {
+        use super::*;
+        use crate::fees::op::LowLevelDriveOperation;
+
+        #[test]
+        fn should_return_positive_balance_ignoring_negative_credit_when_balance_nonzero() {
+            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;
+
+            // Add positive balance
+            drive
+                .add_to_identity_balance(
+                    identity.id().to_buffer(),
+                    added_balance,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("should add balance");
+
+            let negative_amount: u64 = 500;
+
+            // Set negative credit (debt)
+            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");
+
+            // When the balance is nonzero, fetch_identity_balance_include_debt
+            // returns just the positive balance. The negative credit is only
+            // consulted when the balance itself is zero.
+            let balance = drive
+                .fetch_identity_balance_include_debt(
+                    identity.id().to_buffer(),
+                    None,
+                    platform_version,
+                )
+                .expect("should not error")
+                .expect("should have balance");
+
+            assert_eq!(
+                balance, added_balance as i64,
+                "should return positive balance when balance > 0, regardless of negative credit"
+            );
+        }
+    }
+
     mod fetch_identity_negative_balance {
         use super::*;
+        use crate::error::Error;
 
         #[test]
-        fn should_error_for_non_existent_identity() {
+        fn should_error_with_grovedb_error_for_non_existent_identity() {
             let drive = setup_drive_with_initial_state_structure(None);
             let platform_version = PlatformVersion::latest();
 
             let mut drive_operations = vec![];
             // For a non-existent identity, the path doesn't exist, so grove_get_raw
-            // returns an error (PathParentLayerNotFound) since the identity subtree
-            // doesn't exist at all.
+            // returns a GroveDB error (PathParentLayerNotFound) since the identity
+            // subtree doesn't exist at all.
             let result = drive.fetch_identity_negative_balance_operations(
                 [0; 32],
                 true,
@@ -208,7 +276,11 @@ mod tests {
                 platform_version,
             );
 
-            assert!(result.is_err());
+            assert!(
+                matches!(result, Err(Error::GroveDB(_))),
+                "expected GroveDB error for non-existent identity path, got: {:?}",
+                result
+            );
         }
 
         #[test]

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

@@ -384,6 +384,49 @@ mod tests {
                 assert_eq!(*balance, identity.balance());
             }
         }
+
+        #[test]
+        fn should_return_only_existing_identities_when_some_are_missing() {
+            let drive = setup_drive_with_initial_state_structure(None);
+            let platform_version = PlatformVersion::latest();
+
+            let identities: Vec<Identity> =
+                Identity::random_identities(2, 3, Some(42), platform_version)
+                    .expect("expected random identities");
+
+            // Only insert the first identity
+            drive
+                .add_new_identity(
+                    identities[0].clone(),
+                    false,
+                    &BlockInfo::default(),
+                    true,
+                    None,
+                    platform_version,
+                )
+                .expect("expected to add identity");
+
+            // Query for both the existing identity and a non-existent one
+            let ids = vec![
+                identities[0].id().to_buffer(),
+                identities[1].id().to_buffer(),
+            ];
+
+            let balances = drive
+                .fetch_identities_balances(&ids, None, platform_version)
+                .expect("should fetch balances");
+
+            // Only the inserted identity should be returned
+            assert_eq!(balances.len(), 1);
+            let balance = balances
+                .get(&identities[0].id().to_buffer())
+                .expect("should have balance for existing identity");
+            assert_eq!(*balance, identities[0].balance());
+            assert!(
+                balances.get(&identities[1].id().to_buffer()).is_none(),
+                "non-existent identity should not appear in results"
+            );
+        }
     }
 
     mod fetch_optional_identities_balances {
@@ -416,10 +459,17 @@ mod tests {
 
             assert_eq!(balances.len(), 2);
             assert_eq!(
-                *balances.get(&identity.id().to_buffer()).unwrap(),
+                *balances
+                    .get(&identity.id().to_buffer())
+                    .expect("existing identity should have an entry in results"),
                 Some(identity.balance())
             );
-            assert_eq!(*balances.get(&[0xff; 32]).unwrap(), None);
+            assert_eq!(
+                *balances
+                    .get(&[0xff; 32])
+                    .expect("non-existent identity should still have an entry in results"),
+                None
+            );
         }
     }
 
@@ -461,9 +511,13 @@ mod tests {
 
             assert_eq!(balances.len(), 5);
 
-            // BTreeMap keys are always sorted, confirming ascending order
-            let keys: Vec<[u8; 32]> = balances.keys().copied().collect();
-            assert!(keys.windows(2).all(|w| w[0] <= w[1]));
+            // Verify that every inserted identity appears with its correct balance
+            for identity in &identities {
+                let balance = balances
+                    .get(&identity.id().to_buffer())
+                    .expect("should contain balance for inserted identity");
+                assert_eq!(*balance, identity.balance());
+            }
         }
 
         #[test]

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

@@ -89,9 +89,10 @@ mod tests {
                 .expect("should not error")
                 .expect("should have nonce");
 
-            // After merging nonce 1, the nonce value should contain 1
-            // The nonce encoding includes missing revision bits, so just check it's nonzero
-            assert!(nonce > 0);
+            // After merging nonce 1, the stored value should be exactly 1.
+            // When merging nonce 1 from an initial value of 0, there are no missing
+            // revisions in between, so no missing-revision bits are set.
+            assert_eq!(nonce, 1);
         }
 
         #[test]

packages/rs-drive/src/drive/identity/insert/add_new_identity/mod.rs

@@ -272,7 +272,9 @@ mod tests {
     }
 
     #[test]
-    fn should_succeed_reinserting_masternode_identity() {
+    fn should_succeed_reinserting_masternode_identity_and_reenable_keys() {
+        use dpp::identity::identity_public_key::accessors::v0::IdentityPublicKeyGettersV0;
+
         let platform_version = PlatformVersion::latest();
         let drive = setup_drive(None);
 
@@ -297,7 +299,37 @@ mod tests {
             )
             .expect("expected to insert identity");
 
-        // Reinserting the same masternode identity should succeed (re-enable keys)
+        // Disable all keys to simulate a masternode being removed
+        let key_ids: Vec<dpp::identity::KeyID> = identity.public_keys().keys().copied().collect();
+        drive
+            .disable_identity_keys(
+                identity.id().to_buffer(),
+                key_ids.clone(),
+                1000, // disable_at timestamp
+                &BlockInfo::default(),
+                true,
+                Some(&transaction),
+                platform_version,
+            )
+            .expect("expected to disable keys");
+
+        // Verify keys are disabled before re-insertion
+        let fetched_keys_before = drive
+            .fetch_all_identity_keys(
+                identity.id().to_buffer(),
+                Some(&transaction),
+                platform_version,
+            )
+            .expect("expected to fetch keys");
+        for key in fetched_keys_before.values() {
+            assert!(
+                key.is_disabled(),
+                "key {} should be disabled before re-insertion",
+                key.id()
+            );
+        }
+
+        // Reinserting the same masternode identity should succeed and re-enable keys
         let result = drive.add_new_identity(
             identity.clone(),
             true,
@@ -309,15 +341,22 @@ mod tests {
 
         assert!(result.is_ok());
 
-        // Verify keys are still present after reinsertion
-        let fetched_keys = drive
+        // Verify keys are re-enabled after reinsertion
+        let fetched_keys_after = drive
             .fetch_all_identity_keys(
                 identity.id().to_buffer(),
                 Some(&transaction),
                 platform_version,
             )
             .expect("expected to fetch keys");
 
-        assert_eq!(fetched_keys.len(), 5);
+        assert_eq!(fetched_keys_after.len(), 5);
+        for key in fetched_keys_after.values() {
+            assert!(
+                !key.is_disabled(),
+                "key {} should be re-enabled after masternode re-insertion",
+                key.id()
+            );
+        }
     }
 }