From 56ff345a9923a4172c47be74ecd35ba8e99f45bf Mon Sep 17 00:00:00 2001
From: Ed Hennis <ed@ripple.com>
Date: Thu, 17 Apr 2025 16:44:00 -0400
Subject: [PATCH] Add Asset and MPTIssue support to more jtx objects /
 functions

- Unfortunately, to work around some ambiguous symbol compilation
  errors, I had to change the implicit conversion from IOU to Asset to
  a conversion from IOU to PrettyAsset, and add a more explicit
  `asset()` function. This workaround only required changing two
  existing tests, so seems acceptable.
---
 src/test/app/AMMCalc_test.cpp |  2 +-
 src/test/jtx/Env.h            |  6 ++++++
 src/test/jtx/amount.h         | 15 +++++++++++++--
 src/test/jtx/impl/Env.cpp     | 20 ++++++++++++++++++++
 4 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/src/test/app/AMMCalc_test.cpp b/src/test/app/AMMCalc_test.cpp
index bebf2844b6..7349b38766 100644
--- a/src/test/app/AMMCalc_test.cpp
+++ b/src/test/app/AMMCalc_test.cpp
@@ -67,7 +67,7 @@ class AMMCalc_test : public beast::unit_test::suite
             // drops
             else if (match[1] == "XRPA")
                 return XRPAmount{std::stoll(match[2])};
-            return amountFromString(gw[match[1]], match[2]);
+            return amountFromString(gw[match[1]].asset(), match[2]);
         }
         return std::nullopt;
     }
diff --git a/src/test/jtx/Env.h b/src/test/jtx/Env.h
index 6484b2ff1b..f73e096845 100644
--- a/src/test/jtx/Env.h
+++ b/src/test/jtx/Env.h
@@ -469,9 +469,15 @@ public:
         Returns 0 if the trust line does not exist.
     */
     // VFALCO NOTE This should return a unit-less amount
+    PrettyAmount
+    balance(Account const& account, Asset const& asset) const;
+
     PrettyAmount
     balance(Account const& account, Issue const& issue) const;
 
+    PrettyAmount
+    balance(Account const& account, MPTIssue const& mptIssue) const;
+
     /** Return the number of objects owned by an account.
      * Returns 0 if the account does not exist.
      */
diff --git a/src/test/jtx/amount.h b/src/test/jtx/amount.h
index 241b83b6c1..c7510ea9e0 100644
--- a/src/test/jtx/amount.h
+++ b/src/test/jtx/amount.h
@@ -143,6 +143,12 @@ public:
         return amount_;
     }
 
+    inline int
+    signum() const
+    {
+        return amount_.signum();
+    }
+
     operator STAmount const&() const
     {
         return amount_;
@@ -367,6 +373,11 @@ public:
     {
         return {currency, account.id()};
     }
+    Asset
+    asset() const
+    {
+        return issue();
+    }
 
     /** Implicit conversion to Issue or Asset.
 
@@ -377,9 +388,9 @@ public:
     {
         return issue();
     }
-    operator Asset() const
+    operator PrettyAsset() const
     {
-        return issue();
+        return asset();
     }
 
     template <
diff --git a/src/test/jtx/impl/Env.cpp b/src/test/jtx/impl/Env.cpp
index 286e15356a..2b9394cfd9 100644
--- a/src/test/jtx/impl/Env.cpp
+++ b/src/test/jtx/impl/Env.cpp
@@ -199,6 +199,26 @@ Env::balance(Account const& account, Issue const& issue) const
     return {amount, lookup(issue.account).name()};
 }
 
+PrettyAmount
+Env::balance(Account const& account, MPTIssue const& mptIssue) const
+{
+    auto const sle = le(keylet::mptoken(mptIssue.getMptID(), account));
+    if (!sle)
+    {
+        return {STAmount(mptIssue, 0), account.name()};
+    }
+    STAmount const amount{mptIssue, sle->getFieldU64(sfMPTAmount)};
+    return {amount, lookup(mptIssue.getIssuer()).name()};
+}
+
+PrettyAmount
+Env::balance(Account const& account, Asset const& asset) const
+{
+    return std::visit(
+        [&](auto const& issue) { return balance(account, issue); },
+        asset.value());
+}
+
 std::uint32_t
 Env::ownerCount(Account const& account) const
 {