Small improvements and checking of return collateral address

cardano_node_tests/tests/test_plutus_v2_spend_build.py

@@ -2346,9 +2346,6 @@ def test_with_total_return_collateral(
             collateral_amount=amount_for_collateral,
         )
 
-        payment_init_balance = cluster.get_address_balance(payment_addr.address)
-        dst_init_balance = cluster.get_address_balance(dst_addr.address)
-
         #  spend the "locked" UTxO
 
         return_collateral_txouts = [
@@ -2367,36 +2364,42 @@ def test_with_total_return_collateral(
             return_collateral_txouts=return_collateral_txouts if use_return_collateral else (),
         )
 
-        # check that the right amount of collateral was taken
-        dst_balance = cluster.get_address_balance(dst_addr.address)
-        assert (
-            dst_balance == dst_init_balance - redeem_cost.collateral
-        ), f"Collateral was NOT spent from `{dst_addr.address}` correctly"
+        # check that collateral was taken
+        assert not cluster.get_utxo(utxo=collateral_utxos), "Collateral was NOT spent"
 
-        if use_return_collateral or not use_total_collateral:
-            txid_redeem = cluster.get_txid(tx_body_file=tx_output_redeem.out_file)
-            return_col_utxos = cluster.get_utxo(
-                txin=f"{txid_redeem}#{len(tx_output_redeem.txouts) + 1}"
-            )
-            assert return_col_utxos, "Return collateral UTxO was not created"
-
-            change_amount = clusterlib.calculate_utxos_balance(utxos=return_col_utxos)
-
-            if use_return_collateral:
-                assert (
-                    change_amount == return_collateral_amount
-                ), f"Incorrect balance for collateral return address `{dst_addr.address}`"
-            else:
-                # check that the collateral amount charged is according to 'collateralPercentage'
-                collateral_charged = amount_for_collateral - change_amount
-                assert collateral_charged == round(
-                    tx_output_redeem.fee * protocol_params["collateralPercentage"] / 100
-                ), "The collateral amount charged is not the expected"
-
-                payment_balance = cluster.get_address_balance(payment_addr.address)
-                assert payment_balance == payment_init_balance + clusterlib.calculate_utxos_balance(
-                    utxos=return_col_utxos
-                ), f"Incorrect balance for change address `{payment_addr.address}`"
+        txid_redeem = cluster.get_txid(tx_body_file=tx_output_redeem.out_file)
+        return_col_utxos = cluster.get_utxo(
+            txin=f"{txid_redeem}#{len(tx_output_redeem.txouts) + 1}"
+        )
+
+        # when total collateral amount is specified, it is necessary to specify also return
+        # collateral `TxOut`
+        if use_total_collateral and not use_return_collateral:
+            assert not return_col_utxos, "Return collateral UTxO was unexpectedly created"
+            return
+
+        # check that correct return collateral UTxO was created
+        assert return_col_utxos, "Return collateral UTxO was NOT created"
+
+        returned_collateral_amount = clusterlib.calculate_utxos_balance(utxos=return_col_utxos)
+
+        if use_return_collateral:
+            assert (
+                returned_collateral_amount == return_collateral_amount
+            ), f"Incorrect balance for collateral return address `{dst_addr.address}`"
+            assert (
+                return_collateral_txouts[0].address == return_col_utxos[0].address
+            ), "Return collateral address doesn't match"
+        else:
+            # check that the collateral amount charged corresponds to 'collateralPercentage'
+            collateral_charged = amount_for_collateral - returned_collateral_amount
+            assert collateral_charged == round(
+                tx_output_redeem.fee * protocol_params["collateralPercentage"] / 100
+            ), "The collateral amount charged is not the expected amount"
+
+            assert (
+                payment_addr.address == return_col_utxos[0].address
+            ), "Return collateral address doesn't match"
 
     @allure.link(helpers.get_vcs_link())
     def test_collateral_with_tokens(