Merge pull request #1119 from lightninglabs/asset-balance-alignment

Add 'include_leased' flag to ListBalances

Author: Roasbeef

cmd/tapcli/assets.go

@@ -48,6 +48,7 @@ var (
 	assetShowWitnessName         = "show_witness"
 	assetShowSpentName           = "show_spent"
 	assetShowLeasedName          = "show_leased"
+	assetIncludeLeasedName       = "include_leased"
 	assetShowUnconfMintsName     = "show_unconfirmed_mints"
 	assetGroupKeyName            = "group_key"
 	assetGroupAnchorName         = "group_anchor"
@@ -642,6 +643,10 @@ var listAssetBalancesCommand = cli.Command{
 			Name:  groupByGroupName,
 			Usage: "Group asset balances by group key",
 		},
+		cli.BoolFlag{
+			Name:  assetIncludeLeasedName,
+			Usage: "Include leased assets in balances",
+		},
 		cli.StringFlag{
 			Name: assetIDName,
 			Usage: "A specific asset ID to run the balance query " +
@@ -663,7 +668,9 @@ func listAssetBalances(ctx *cli.Context) error {
 
 	var err error
 
-	req := &taprpc.ListBalancesRequest{}
+	req := &taprpc.ListBalancesRequest{
+		IncludeLeased: ctx.Bool(assetIncludeLeasedName),
+	}
 
 	if !ctx.Bool(groupByGroupName) {
 		req.GroupBy = &taprpc.ListBalancesRequest_AssetId{

itest/assertions.go

@@ -1886,7 +1886,7 @@ func AssertGenesisOutput(t *testing.T, output *taprpc.ManagedUtxo,
 }
 
 func AssertAssetBalances(t *testing.T, client taprpc.TaprootAssetsClient,
-	simpleAssets, issuableAssets []*taprpc.Asset) {
+	simpleAssets, issuableAssets []*taprpc.Asset, includeLeased bool) {
 
 	t.Helper()
 
@@ -1901,7 +1901,8 @@ func AssertAssetBalances(t *testing.T, client taprpc.TaprootAssetsClient,
 	}
 	assetIDBalances, err := client.ListBalances(
 		ctxt, &taprpc.ListBalancesRequest{
-			GroupBy: balanceReq,
+			GroupBy:       balanceReq,
+			IncludeLeased: includeLeased,
 		},
 	)
 	require.NoError(t, err)
@@ -1912,7 +1913,9 @@ func AssertAssetBalances(t *testing.T, client taprpc.TaprootAssetsClient,
 
 	require.Equal(t, len(allAssets), len(assetIDBalances.AssetBalances))
 
+	var totalBalance uint64
 	for _, balance := range assetIDBalances.AssetBalances {
+		totalBalance += balance.Balance
 		for _, rpcAsset := range allAssets {
 			balanceGen := balance.AssetGenesis
 			targetGen := rpcAsset.AssetGenesis
@@ -1931,6 +1934,20 @@ func AssertAssetBalances(t *testing.T, client taprpc.TaprootAssetsClient,
 		}
 	}
 
+	// We should also ensure that the total balance returned by
+	// `ListBalances` matches the total balance returned by `ListAssets`.
+	assetList, err := client.ListAssets(ctxt, &taprpc.ListAssetRequest{
+		IncludeLeased: includeLeased,
+	})
+	require.NoError(t, err)
+
+	var totalAssetListBalance uint64
+	for _, asset := range assetList.Assets {
+		totalAssetListBalance += asset.Amount
+	}
+
+	require.Equal(t, totalBalance, totalAssetListBalance)
+
 	// We'll also ensure that we're able to get the balance by key group
 	// for all the assets that have one specified.
 	groupBalanceReq := &taprpc.ListBalancesRequest_GroupKey{

itest/assets_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/lightninglabs/taproot-assets/internal/test"
 	"github.com/lightninglabs/taproot-assets/proof"
 	"github.com/lightninglabs/taproot-assets/taprpc"
+	wrpc "github.com/lightninglabs/taproot-assets/taprpc/assetwalletrpc"
 	"github.com/lightninglabs/taproot-assets/taprpc/mintrpc"
 	"github.com/lightninglabs/taproot-assets/taprpc/tapdevrpc"
 	"github.com/lightningnetwork/lnd/lnrpc"
@@ -109,7 +110,9 @@ func testMintAssets(t *harnessTest) {
 	// Now that all our assets have been issued, we'll use the balance
 	// calls to ensure that we're able to retrieve the proper balance for
 	// them all.
-	AssertAssetBalances(t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets)
+	AssertAssetBalances(
+		t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets, false,
+	)
 
 	// Check that we can retrieve the group keys for the issuable assets.
 	assertGroups(t.t, t.tapd, issuableAssets)
@@ -441,7 +444,9 @@ func testMintAssetsWithTapscriptSibling(t *harnessTest) {
 	rpcIssuableAssets := MintAssetsConfirmBatch(
 		t.t, t.lndHarness.Miner.Client, t.tapd, issuableAssets,
 	)
-	AssertAssetBalances(t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets)
+	AssertAssetBalances(
+		t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets, false,
+	)
 
 	// Filter the managed UTXOs to select the genesis UTXO with the
 	// tapscript sibling.
@@ -617,3 +622,77 @@ func testMintBatchAndTransfer(t *harnessTest) {
 
 	require.True(t.t, proto.Equal(originalBatch, afterBatch))
 }
+
+// testAssetBalances tests the balance retrieval functionality for issued
+// assets. The function mints two batches of assets and asserts if the tapcli
+// `assets balance` returns the correct balances. It then funds a vPSBT, putting
+// a lease on one of the two batches. It then asserts whether the endpoint still
+// returns the correct balances, taking into account the `include_leased` flag.
+func testAssetBalances(t *harnessTest) {
+	ctxb := context.Background()
+	ctxt, cancel := context.WithTimeout(ctxb, defaultWaitTimeout)
+	defer cancel()
+
+	rpcSimpleAssets := MintAssetsConfirmBatch(
+		t.t, t.lndHarness.Miner.Client, t.tapd, simpleAssets,
+	)
+	rpcIssuableAssets := MintAssetsConfirmBatch(
+		t.t, t.lndHarness.Miner.Client, t.tapd, issuableAssets,
+	)
+	targetAsset := rpcSimpleAssets[0]
+
+	// Now that all our assets have been issued, we'll use the balance
+	// calls to ensure that we're able to retrieve the proper balance for
+	// them all.
+	AssertAssetBalances(
+		t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets, false,
+	)
+
+	var (
+		targetAssetGenesis = targetAsset.AssetGenesis
+		aliceTapd          = t.tapd
+		bobLnd             = t.lndHarness.Bob
+	)
+
+	// We create a second tapd node that will be used to simulate a second
+	// party in the test. This tapd node is connected to lnd "Bob".
+	bobTapd := setupTapdHarness(t.t, t, bobLnd, t.universeServer)
+	defer func() {
+		require.NoError(t.t, bobTapd.stop(!*noDelete))
+	}()
+
+	const assetsToSend = 1000
+	bobAddr, err := bobTapd.NewAddr(ctxt, &taprpc.NewAddrRequest{
+		AssetId: targetAssetGenesis.AssetId,
+		Amt:     assetsToSend,
+	})
+	require.NoError(t.t, err)
+
+	// Now we can create our virtual transaction and ask Alice's tapd to
+	// fund it.
+	recipients := map[string]uint64{
+		bobAddr.Encoded: bobAddr.Amount,
+	}
+	_, err = aliceTapd.FundVirtualPsbt(
+		ctxt, &wrpc.FundVirtualPsbtRequest{
+			Template: &wrpc.FundVirtualPsbtRequest_Raw{
+				Raw: &wrpc.TxTemplate{
+					Recipients: recipients,
+				},
+			},
+		},
+	)
+	require.NoError(t.t, err)
+
+	// With a transaction funding should have led to a lease on the simple
+	// assets, we'll use the balance calls to ensure that we're able to
+	// retrieve the proper balances.
+	rpcEmptyAssets := []*taprpc.Asset{}
+	AssertAssetBalances(
+		t.t, t.tapd, rpcEmptyAssets, rpcIssuableAssets, false,
+	)
+
+	AssertAssetBalances(
+		t.t, t.tapd, rpcSimpleAssets, rpcIssuableAssets, true,
+	)
+}

itest/test_list_on_test.go

@@ -17,6 +17,10 @@ var testCases = []*testCase{
 		name: "mint batch and transfer",
 		test: testMintBatchAndTransfer,
 	},
+	{
+		name: "asset balances",
+		test: testAssetBalances,
+	},
 	{
 		name: "asset meta validation",
 		test: testAssetMeta,

rpcserver.go

@@ -922,7 +922,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case assetID != nil:
 		// Retrieve the current asset balance.
 		balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
-			ctx, assetID,
+			ctx, assetID, true,
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query asset balance: %w",
@@ -938,7 +938,7 @@ func (r *rpcServer) checkBalanceOverflow(ctx context.Context,
 	case groupPubKey != nil:
 		// Retrieve the current balance of the group.
 		balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-			ctx, groupPubKey,
+			ctx, groupPubKey, true,
 		)
 		if err != nil {
 			return fmt.Errorf("unable to query group balance: %w",
@@ -1106,9 +1106,12 @@ func (r *rpcServer) MarshalChainAsset(ctx context.Context, a *asset.ChainAsset,
 }
 
 func (r *rpcServer) listBalancesByAsset(ctx context.Context,
-	assetID *asset.ID) (*taprpc.ListBalancesResponse, error) {
+	assetID *asset.ID, includeLeased bool) (*taprpc.ListBalancesResponse,
+	error) {
 
-	balances, err := r.cfg.AssetStore.QueryBalancesByAsset(ctx, assetID)
+	balances, err := r.cfg.AssetStore.QueryBalancesByAsset(
+		ctx, assetID, includeLeased,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
 	}
@@ -1138,10 +1141,11 @@ func (r *rpcServer) listBalancesByAsset(ctx context.Context,
 }
 
 func (r *rpcServer) listBalancesByGroupKey(ctx context.Context,
-	groupKey *btcec.PublicKey) (*taprpc.ListBalancesResponse, error) {
+	groupKey *btcec.PublicKey,
+	includeLeased bool) (*taprpc.ListBalancesResponse, error) {
 
 	balances, err := r.cfg.AssetStore.QueryAssetBalancesByGroup(
-		ctx, groupKey,
+		ctx, groupKey, includeLeased,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("unable to list balances: %w", err)
@@ -1293,7 +1297,7 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 			copy(assetID[:], req.AssetFilter)
 		}
 
-		return r.listBalancesByAsset(ctx, assetID)
+		return r.listBalancesByAsset(ctx, assetID, req.IncludeLeased)
 
 	case *taprpc.ListBalancesRequest_GroupKey:
 		if !groupBy.GroupKey {
@@ -1310,7 +1314,9 @@ func (r *rpcServer) ListBalances(ctx context.Context,
 			}
 		}
 
-		return r.listBalancesByGroupKey(ctx, groupKey)
+		return r.listBalancesByGroupKey(
+			ctx, groupKey, req.IncludeLeased,
+		)
 
 	default:
 		return nil, fmt.Errorf("invalid group_by")