From 146bf5102197852587343759f9d6b2b2369507d2 Mon Sep 17 00:00:00 2001
From: sowle <crypto.sowle@gmail.com>
Date: Wed, 18 Oct 2023 21:08:16 +0200
Subject: [PATCH] chaingen: balance_via_wallet() now supports asset id,
 replace_coinbase_in_genesis_block() fixed

---
 tests/core_tests/chaingen.cpp | 34 ++++++++++++++++++++++++----------
 tests/core_tests/chaingen.h   |  5 +++--
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/tests/core_tests/chaingen.cpp b/tests/core_tests/chaingen.cpp
index 2f0aac77..b16e24ee 100644
--- a/tests/core_tests/chaingen.cpp
+++ b/tests/core_tests/chaingen.cpp
@@ -1905,10 +1905,10 @@ bool find_block_chain(const std::vector<test_event_entry>& events, std::vector<c
   return b_success;
 }
 
-void balance_via_wallet(const tools::wallet2& w, uint64_t* p_total, uint64_t* p_unlocked, uint64_t* p_awaiting_in, uint64_t* p_awaiting_out, uint64_t* p_mined)
+void balance_via_wallet(const tools::wallet2& w, const crypto::public_key& asset_id, uint64_t* p_total, uint64_t* p_unlocked, uint64_t* p_awaiting_in, uint64_t* p_awaiting_out, uint64_t* p_mined)
 {
   uint64_t total, unlocked, awaiting_in, awaiting_out, mined;
-  total = w.balance(unlocked, awaiting_in, awaiting_out, mined);
+  total = w.balance(unlocked, awaiting_in, awaiting_out, mined, asset_id);
 
   if (p_total)
     *p_total = total;
@@ -1923,12 +1923,16 @@ void balance_via_wallet(const tools::wallet2& w, uint64_t* p_total, uint64_t* p_
 }
 
 bool check_balance_via_wallet(const tools::wallet2& w, const char* account_name,
-  uint64_t expected_total, uint64_t expected_mined, uint64_t expected_unlocked, uint64_t expected_awaiting_in, uint64_t expected_awaiting_out)
+  uint64_t expected_total, uint64_t expected_mined, uint64_t expected_unlocked, uint64_t expected_awaiting_in, uint64_t expected_awaiting_out, const crypto::public_key& asset_id /* = currency::native_coin_asset_id */)
 {
   uint64_t total, unlocked, awaiting_in, awaiting_out, mined;
-  balance_via_wallet(w, &total, &unlocked, &awaiting_in, &awaiting_out, &mined);
+  balance_via_wallet(w, asset_id, &total, &unlocked, &awaiting_in, &awaiting_out, &mined);
 
-  LOG_PRINT_CYAN("Balance for wallet " << account_name << " @ height " << w.get_top_block_height() << ":" << ENDL <<
+  std::string asset_id_str;
+  if (asset_id != currency::native_coin_asset_id)
+    asset_id_str = std::string(", asset_id: ") + epee::string_tools::pod_to_hex(asset_id).erase(4, 56).insert(4, "...");
+
+  LOG_PRINT_CYAN("Balance for wallet " << account_name << " @ height " << w.get_top_block_height() << asset_id_str << ":" << ENDL <<
     "unlocked:     " << print_money(unlocked) << ENDL <<
     "awaiting in:  " << print_money(awaiting_in) << ENDL <<
     "awaiting out: " << print_money(awaiting_out) << ENDL <<
@@ -1939,7 +1943,7 @@ bool check_balance_via_wallet(const tools::wallet2& w, const char* account_name,
 
   bool r = true;
 
-#define _CHECK_BAL(v) if (!(expected_##v == INVALID_BALANCE_VAL || v == expected_##v)) { r = false; LOG_PRINT_RED_L0("invalid " #v " balance, expected: " << print_money_brief(expected_##v)); }
+#define _CHECK_BAL(v) if (!(expected_##v == INVALID_BALANCE_VAL || v == expected_##v)) { r = false; LOG_PRINT_RED_L0("invalid " #v " balance, expected: " << print_money_brief(expected_##v) << asset_id_str); }
   _CHECK_BAL(unlocked)
   _CHECK_BAL(awaiting_in)
   _CHECK_BAL(awaiting_out)
@@ -2329,6 +2333,7 @@ bool shuffle_source_entries(std::vector<tx_source_entry>& sources)
   return true;
 }
 
+// creates destinations.size() + 1 outputs if the total sum of amounts is less than the original premine amount (the last one will have amount = old_premine - sum)
 bool replace_coinbase_in_genesis_block(const std::vector<currency::tx_destination_entry>& destinations, test_generator& generator, std::vector<test_event_entry>& events, currency::block& genesis_block)
 {
   bool r = false;
@@ -2340,13 +2345,19 @@ bool replace_coinbase_in_genesis_block(const std::vector<currency::tx_destinatio
 
   // replace tx key
   keypair tx_key = keypair::generate();
-  for(auto& el : genesis_block.miner_tx.extra)
+  for(auto it = genesis_block.miner_tx.extra.begin(); it != genesis_block.miner_tx.extra.end(); /* nothing */)
   {
-    if (el.type() == typeid(crypto::public_key))
+    if (it->type() == typeid(crypto::public_key))
     {
-      boost::get<crypto::public_key>(el) = tx_key.pub;
-      break;
+      boost::get<crypto::public_key>(*it) = tx_key.pub;    // rewtire it
     }
+    else if (it->type() == typeid(tx_derivation_hint))
+    {
+      it = genesis_block.miner_tx.extra.erase(it);
+      continue;                                            // remove it
+    }
+
+    ++it;
   }
   uint64_t total_amount = 0;
 
@@ -2357,10 +2368,13 @@ bool replace_coinbase_in_genesis_block(const std::vector<currency::tx_destinatio
   {
     uint64_t amount = output_index < destinations.size() ? destinations[output_index].amount : premine_amount - total_amount;
     const account_public_address& addr = output_index < destinations.size() ? destinations[output_index].addr.back() : destinations.back().addr.back();
+    if (amount == 0)
+      break;
 
     crypto::key_derivation derivation{};
     bool r = crypto::generate_key_derivation(addr.view_public_key, tx_key.sec, derivation);
     CHECK_AND_ASSERT_MES(r, false, "generate_key_derivation failed");
+    genesis_block.miner_tx.extra.emplace_back(make_tx_derivation_hint_from_uint16(get_derivation_hint(derivation)));
 
     txout_to_key target{};
     r = crypto::derive_public_key(derivation, output_index, addr.spend_public_key, target.key);
diff --git a/tests/core_tests/chaingen.h b/tests/core_tests/chaingen.h
index 69f1ab56..e1029aae 100644
--- a/tests/core_tests/chaingen.h
+++ b/tests/core_tests/chaingen.h
@@ -703,14 +703,15 @@ bool fill_tx_sources_and_destinations(const std::vector<test_event_entry>& event
                                       bool use_ref_by_id = false);
 uint64_t get_balance(const currency::account_keys& addr, const std::vector<currency::block>& blockchain, const map_hash2tx_t& mtx, bool dbg_log = false);
 uint64_t get_balance(const currency::account_base& addr, const std::vector<currency::block>& blockchain, const map_hash2tx_t& mtx, bool dbg_log = false);
-void balance_via_wallet(const tools::wallet2& w, uint64_t* p_total, uint64_t* p_unlocked = 0, uint64_t* p_awaiting_in = 0, uint64_t* p_awaiting_out = 0, uint64_t* p_mined = 0);
+void balance_via_wallet(const tools::wallet2& w, const crypto::public_key& asset_id, uint64_t* p_total, uint64_t* p_unlocked = 0, uint64_t* p_awaiting_in = 0, uint64_t* p_awaiting_out = 0, uint64_t* p_mined = 0);
 #define INVALID_BALANCE_VAL std::numeric_limits<uint64_t>::max()
 bool check_balance_via_wallet(const tools::wallet2& w, const char* account_name,
                               uint64_t expected_total,
                               uint64_t expected_mined = INVALID_BALANCE_VAL,
                               uint64_t expected_unlocked = INVALID_BALANCE_VAL,
                               uint64_t expected_awaiting_in = INVALID_BALANCE_VAL,
-                              uint64_t expected_awaiting_out = INVALID_BALANCE_VAL);
+                              uint64_t expected_awaiting_out = INVALID_BALANCE_VAL,
+                              const crypto::public_key& asset_id = currency::native_coin_asset_id);
 
 bool calculate_amounts_many_outs_have_and_no_outs_have(const uint64_t first_blocks_reward, uint64_t& amount_many_outs_have, uint64_t& amount_no_outs_have);
 bool find_global_index_for_output(const std::vector<test_event_entry>& events, const crypto::hash& head_block_hash, const currency::transaction& reference_tx, const size_t reference_tx_out_index, uint64_t& global_index);