Add the upgrade test

crates/sui-transaction-builder/src/lib.rs

@@ -493,12 +493,14 @@ mod tests {
     use sui_types::types::Identifier;
     use sui_types::types::ObjectDigest;
     use sui_types::types::ObjectId;
+    use sui_types::types::ObjectType;
     use sui_types::types::Transaction;
     use sui_types::types::TransactionEffects;
     use sui_types::types::TypeTag;
 
     use crate::object::Object;
     use crate::Function;
+    use crate::RawBytes;
     use crate::Serialized;
     use crate::TransactionBuilder;
 
@@ -883,28 +885,22 @@ mod tests {
         let sig = pk.sign_transaction(&tx).unwrap();
         let effects = client.execute_tx(vec![sig], &tx).await;
         let mut package_id: Option<ObjectId> = None;
-        let mut obj_digest: Option<ObjectDigest> = None;
+        let mut created_objs = vec![];
         if let Ok(Some(ref effects)) = effects {
             match effects {
                 TransactionEffects::V2(e) => {
                     println!("Tx digest: {:?}", e.transaction_digest);
 
                     for obj in e.changed_objects.clone() {
-                        println!("Obj ID: {}", obj.object_id);
-                        println!("Change type: {:?}", obj.change.id_operation);
-                        println!("Obj output state: {:?}", obj.change.output_state);
                         match obj.change.id_operation {
                             IdOperation::Created => {
                                 let change = obj.change.output_state;
                                 match change {
-                                    sui_types::types::ObjectOut::PackageWrite {
-                                        version,
-                                        digest,
-                                    } => {
+                                    sui_types::types::ObjectOut::PackageWrite { .. } => {
                                         package_id = Some(obj.object_id);
-                                        obj_digest = Some(digest);
-                                        println!("Package version: {:?}", version);
-                                        println!("Package digest: {:?}", digest);
+                                    }
+                                    sui_types::types::ObjectOut::ObjectWrite { .. } => {
+                                        created_objs.push(obj.object_id);
                                     }
                                     _ => {}
                                 }
@@ -918,7 +914,7 @@ mod tests {
         }
         wait_for_tx_and_check_effects_status_success(&client, &tx, effects).await;
 
-        let modules: [u8; 587] = [
+        let updated_modules: [u8; 587] = [
             161, 28, 235, 11, 6, 0, 0, 0, 10, 1, 0, 8, 2, 8, 16, 3, 24, 51, 4, 75, 2, 5, 77, 57, 7,
             134, 1, 155, 1, 8, 161, 2, 64, 10, 225, 2, 18, 12, 243, 2, 163, 1, 13, 150, 4, 6, 0, 4,
             1, 11, 1, 16, 1, 17, 0, 1, 12, 0, 0, 0, 8, 0, 1, 3, 4, 0, 3, 2, 2, 0, 0, 6, 0, 1, 0, 0,
@@ -945,28 +941,67 @@ mod tests {
             20, 6, 1, 0, 0, 0, 0, 0, 0, 0, 22, 11, 0, 15, 2, 21, 11, 3, 17, 7, 11, 1, 11, 2, 18, 0,
             2, 0, 1, 0, 2, 1, 1, 0,
         ];
+
         let mut tx = TransactionBuilder::new();
-        let sender = tx.input(Serialized(&address));
-        let upgrade_arg = tx.input(Serialized(&0));
-        let digest_arg = tx.input(Serialized(&obj_digest.unwrap()));
 
+        for o in created_objs {
+            let obj = client.object(*o.as_address(), None).await.unwrap().unwrap();
+            match obj.object_type() {
+                ObjectType::Struct(x) if x.name.to_string() == "UpgradeCap" => {
+                    match obj.owner() {
+                        sui_types::types::Owner::Address(_) => {
+                            tx.input(Object::immutable(o, obj.version(), obj.digest()))
+                        }
+                        sui_types::types::Owner::Shared(x) => tx.input(Object::shared(o, *x, true)),
+                        // If the capability is owned by an object, then the module defining the owning
+                        // object gets to decide how the upgrade capability should be used.
+                        sui_types::types::Owner::Object(_) => {
+                            panic!("Upgrade capability controlled by object")
+                        }
+                        sui_types::types::Owner::Immutable => panic!("Upgrade capability is stored immutably and cannot be used for upgrades"),
+                    };
+                    break;
+                }
+                _ => {}
+            };
+        }
+
+        let upgrade_policy = tx.input(Serialized(&0u8));
+        let package_digest = [
+            68, 89, 156, 51, 190, 35, 155, 216, 248, 49, 135, 170, 106, 42, 190, 4, 208, 59, 155,
+            89, 74, 63, 70, 95, 207, 78, 227, 22, 136, 146, 57, 79,
+        ];
+        let digest_arg = tx.input(RawBytes(package_digest.to_vec()));
+
+        // we need this ticket to authorize the upgrade
         let upgrade_ticket = tx.move_call(
             Function::new(
-                "0x1".parse().unwrap(),
+                "0x2".parse().unwrap(),
                 Identifier::from_str("package").unwrap(),
                 Identifier::from_str("authorize_upgrade").unwrap(),
                 vec![],
             ),
-            vec![Argument::Input(0), upgrade_arg, digest_arg],
+            vec![Argument::Input(0), upgrade_policy, digest_arg],
         );
-
-        let upgrade_cap = tx.upgrade(
-            vec![modules.to_vec()],
+        // now we can upgrade the package
+        let upgrade_receipt = tx.upgrade(
+            vec![updated_modules.to_vec()],
             deps,
             package_id.unwrap(),
             upgrade_ticket,
         );
 
+        // commit the upgrade
+        tx.move_call(
+            Function::new(
+                "0x2".parse().unwrap(),
+                Identifier::from_str("package").unwrap(),
+                Identifier::from_str("commit_upgrade").unwrap(),
+                vec![],
+            ),
+            vec![Argument::Input(0), upgrade_receipt],
+        );
+
         let gas = coins.last().unwrap().id;
         let gas_obj = client.object(gas.into(), None).await.unwrap().unwrap();
         tx.set_gas(vec![Object::owned(
@@ -977,11 +1012,11 @@ mod tests {
         tx.set_gas_budget(500000000);
         tx.set_gas_price(1000);
         tx.set_sender(address);
+        let sender = tx.input(Serialized(&address));
         tx.transfer_objects(vec![upgrade_cap], sender);
         let tx = tx.finish().unwrap();
         let sig = pk.sign_transaction(&tx).unwrap();
         let effects = client.execute_tx(vec![sig], &tx).await;
-        println!("Upgrade effects {:?}", effects);
         wait_for_tx_and_check_effects_status_success(&client, &tx, effects).await;
     }
 }