[move][testing] add test_scenario::take_shared_by_id, `can_take_sha…

…red_by_id`, `can_take_shared`

This brings the shared object functionality for `test_scenario` up to par with the owned object functionality.

Added tests for the new functions + `take_shared`, which was not tested

crates/sui-framework/sources/test_scenario.move

@@ -195,6 +195,56 @@ module sui::test_scenario {
         }
     }
 
+    /// Same as `take_shared`, but returns the object of type `T` with object ID `id`.
+    /// Should only be used in cases where current tx sender has more than one shared object of
+    /// type `T` in their inventory.
+    /// Returns a wrapper that only supports a `borrow_mut` API to get the mutable reference.
+    public fun take_shared_by_id<T: key>(scenario: &mut Scenario, id: ID): SharedWrapper<T> {
+        let objects: vector<T> = get_unowned_inventory<T>(
+            false, /* immutable */
+            last_tx_start_index(scenario)
+        );
+        let object_opt = find_object_by_id_in_inventory(objects, &id);
+        let object = remove_unique_object_from_inventory(scenario, option::to_vec(object_opt));
+        SharedWrapper {
+            object,
+        }
+    }
+
+     /// Return `true` if a call to `take_shared<T>(scenario)` will succeed
+    public fun can_take_shared<T: key>(scenario: &Scenario): bool {
+        let objects: vector<T> = get_unowned_inventory<T>(
+            false, /* immutable */
+            last_tx_start_index(scenario)
+        );
+        // Check that there is one unique such object, and it has not
+        // yet been removed from the inventory.
+        let res = vector::length(&objects) == 1;
+        if (res) {
+            let id = object::borrow_id(vector::borrow(&objects, 0));
+            res = !vector::contains(&scenario.removed, id);
+        };
+        drop_object_for_testing(objects);
+        res
+    }
+
+    /// This function tells you whether calling `take_shared_by_id` would succeed.
+    /// It provides a way to check without triggering assertions.
+    public fun can_take_shared_by_id<T: key>(scenario: &Scenario, id: ID): bool {
+        // Check that the object has not been removed from the inventory.
+        if (vector::contains(&scenario.removed, &id)) {
+            return false
+        };
+        let objects: vector<T> = get_unowned_inventory<T>(
+            false, /* immutable */
+            last_tx_start_index(scenario)
+        );
+        let object_opt: Option<T> = find_object_by_id_in_inventory(objects, &id);
+        let res =  option::is_some(&object_opt);
+        drop_object_for_testing(object_opt);
+        res
+    }
+
     /// Remove the object of type `T` from the shared inventory that wast most recently created.
     /// Aborts if there is no object of type `T` in the inventory.
     public fun take_last_created_shared<T: key>(scenario: &mut Scenario): SharedWrapper<T> {

crates/sui-framework/tests/test_scenario_tests.move

@@ -381,6 +381,68 @@ module sui::test_scenarioTests {
             };
     }
 
+    #[test]
+    fun test_take_shared_by_id() {
+        let sender = @0x0;
+        let scenario = test_scenario::begin(&sender);
+        let uid1 = test_scenario::new_object(&mut scenario);
+        let uid2 = test_scenario::new_object(&mut scenario);
+        let uid3 = test_scenario::new_object(&mut scenario);
+        let id1 = object::uid_to_inner(&uid1);
+        let id2 = object::uid_to_inner(&uid2);
+        let id3 = object::uid_to_inner(&uid3);
+        {
+            let obj1 = Object { id: uid1, value: 10 };
+            let obj2 = Object { id: uid2, value: 20 };
+            let obj3 = Object { id: uid3, value: 30 };
+            transfer::share_object(obj1);
+            transfer::share_object(obj2);
+            transfer::share_object(obj3)
+        };
+        test_scenario::next_tx(&mut scenario, &sender);
+        {
+            assert!(
+                test_scenario::can_take_shared_by_id<Object>(&scenario, id1),
+                OBJECT_ID_NOT_FOUND
+            );
+            assert!(
+                test_scenario::can_take_shared_by_id<Object>(&scenario, id2),
+                OBJECT_ID_NOT_FOUND
+            );
+            assert!(
+                test_scenario::can_take_shared_by_id<Object>(&scenario, id3),
+                OBJECT_ID_NOT_FOUND
+            );
+            let obj1 = test_scenario::take_shared_by_id<Object>(&mut scenario, id1);
+            let obj3 = test_scenario::take_shared_by_id<Object>(&mut scenario, id3);
+            let obj2 = test_scenario::take_shared_by_id<Object>(&mut scenario, id2);
+            assert!(test_scenario::borrow_mut(&mut obj1).value == 10, VALUE_MISMATCH);
+            assert!(test_scenario::borrow_mut(&mut obj2).value == 20, VALUE_MISMATCH);
+            assert!(test_scenario::borrow_mut(&mut obj3).value == 30, VALUE_MISMATCH);
+            test_scenario::return_shared(&mut scenario, obj1);
+            test_scenario::return_shared(&mut scenario, obj2);
+            test_scenario::return_shared(&mut scenario, obj3);
+        };
+    }
+
+    #[test]
+    fun test_take_shared() {
+        let sender = @0x0;
+        let scenario = test_scenario::begin(&sender);
+        let uid1 = test_scenario::new_object(&mut scenario);
+        {
+            let obj1 = Object { id: uid1, value: 10 };
+            transfer::share_object(obj1);
+        };
+        test_scenario::next_tx(&mut scenario, &sender);
+        {
+            assert!(test_scenario::can_take_shared<Object>(&scenario), 1);
+            let obj1 = test_scenario::take_shared<Object>(&mut scenario);
+            assert!(test_scenario::borrow_mut(&mut obj1).value == 10, VALUE_MISMATCH);
+            test_scenario::return_shared(&mut scenario, obj1);
+        }
+    }
+
     /// Create object and parent. object is a child of parent.
     /// parent is owned by sender of `scenario`.
     fun create_parent_and_object(scenario: &mut Scenario) {