[Neutrino] Update balances (#888)

* [Neutrino] Update balances

* responded to the comments

* some more changes

app/server/src/main/scala/org/bitcoins/server/Main.scala

@@ -77,12 +77,17 @@ object Main extends App {
 
       val callbacks = {
         import DataMessageHandler._
-        val onTX: OnTxReceived = { tx =>
+        val onTx: OnTxReceived = { tx =>
           wallet.processTransaction(tx, confirmations = 0)
           ()
         }
+        val onBlock: OnBlockReceived = { block =>
+          wallet.processBlock(block, 0)
+          ()
+        }
 
-        SpvNodeCallbacks(onTxReceived = Seq(onTX))
+        SpvNodeCallbacks(onTxReceived = Seq(onTx),
+                         onBlockReceived = Seq(onBlock))
       }
       if (nodeConf.isSPVEnabled) {
         for {

node-test/src/test/scala/org/bitcoins/node/NeutrinoNodeWithWalletTest.scala

@@ -0,0 +1,121 @@
+package org.bitcoins.node
+
+import akka.actor.Cancellable
+import org.bitcoins.core.crypto.DoubleSha256DigestBE
+import org.bitcoins.core.currency._
+import org.bitcoins.node.networking.peer.DataMessageHandler
+import org.bitcoins.rpc.client.common.BitcoindVersion
+import org.bitcoins.rpc.util.AsyncUtil
+import org.bitcoins.server.BitcoinSAppConfig
+import org.bitcoins.testkit.BitcoinSTestAppConfig
+import org.bitcoins.testkit.fixtures.UsesExperimentalBitcoind
+import org.bitcoins.testkit.node.NodeUnitTest.NeutrinoNodeFundedWalletBitcoind
+import org.bitcoins.testkit.node.{NodeTestUtil, NodeUnitTest}
+import org.bitcoins.testkit.wallet.BitcoinSWalletTest
+import org.bitcoins.wallet.api.UnlockedWalletApi
+import org.bitcoins.wallet.config.WalletAppConfig
+import org.bitcoins.wallet.models.SpendingInfoTable
+import org.scalatest.exceptions.TestFailedException
+import org.scalatest.{DoNotDiscover, FutureOutcome}
+
+import scala.concurrent.duration._
+import scala.concurrent.{Future, Promise}
+
+@DoNotDiscover
+class NeutrinoNodeWithWalletTest extends NodeUnitTest {
+
+  /** Wallet config with data directory set to user temp directory */
+  implicit override protected def config: BitcoinSAppConfig =
+    BitcoinSTestAppConfig.getNeutrinoTestConfig()
+
+  override type FixtureParam = NeutrinoNodeFundedWalletBitcoind
+
+  def withFixture(test: OneArgAsyncTest): FutureOutcome = {
+    withNeutrinoNodeFundedWalletBitcoind(test,
+                                         callbacks,
+                                         Some(BitcoindVersion.Experimental))
+  }
+
+  private val assertionP: Promise[Boolean] = Promise()
+
+  private val walletP: Promise[UnlockedWalletApi] = Promise()
+  private val walletF: Future[UnlockedWalletApi] = walletP.future
+
+  val amountFromBitcoind = 1.bitcoin
+
+  def callbacks: SpvNodeCallbacks = {
+    val onBlock: DataMessageHandler.OnBlockReceived = { block =>
+      for {
+        wallet <- walletF
+        _ <- wallet.processBlock(block, confirmations = 0)
+      } yield ()
+    }
+    SpvNodeCallbacks(
+      onBlockReceived = Seq(onBlock)
+    )
+  }
+
+  it must "receive information about received payments" taggedAs (UsesExperimentalBitcoind) in {
+    param =>
+      val NeutrinoNodeFundedWalletBitcoind(node, wallet, bitcoind) = param
+
+      walletP.success(wallet)
+
+      def clearSpendingInfoTable(): Future[Int] = {
+        import slick.jdbc.SQLiteProfile.api._
+
+        val conf: WalletAppConfig = wallet.walletConfig
+        val table = TableQuery[SpendingInfoTable]
+        conf.database.run(table.delete)
+      }
+
+      def condition(): Future[Boolean] = {
+        for {
+          balance <- wallet.getUnconfirmedBalance()
+          addresses <- wallet.listAddresses()
+          utxos <- wallet.listUtxos()
+        } yield {
+          balance == BitcoinSWalletTest.initialFunds + amountFromBitcoind &&
+          utxos.size == 2 &&
+          addresses.map(_.scriptPubKey) == utxos.map(_.output.scriptPubKey)
+        }
+      }
+
+      for {
+        addresses <- wallet.listAddresses()
+        utxos <- wallet.listUtxos()
+        _ = assert(addresses.size == 1)
+        _ = assert(utxos.size == 1)
+
+        _ <- node.sync()
+        _ <- NodeTestUtil.awaitSync(node, bitcoind)
+        _ <- NodeTestUtil.awaitCompactFiltersSync(node, bitcoind)
+
+        address <- wallet.getNewAddress()
+        _ <- bitcoind
+          .sendToAddress(address, amountFromBitcoind)
+
+        addresses <- wallet.listAddresses()
+        utxos <- wallet.listUtxos()
+        _ = assert(addresses.size == 2)
+        _ = assert(utxos.size == 1)
+
+        _ <- clearSpendingInfoTable()
+
+        addresses <- wallet.listAddresses()
+        utxos <- wallet.listUtxos()
+        _ = assert(addresses.size == 2)
+        _ = assert(utxos.size == 0)
+
+        _ <- bitcoind.getNewAddress
+          .flatMap(bitcoind.generateToAddress(1, _))
+        _ <- NodeTestUtil.awaitSync(node, bitcoind)
+        _ <- NodeTestUtil.awaitCompactFiltersSync(node, bitcoind)
+
+        addresses <- wallet.listAddresses()
+        _ <- node.rescan(addresses.map(_.scriptPubKey))
+
+        _ <- AsyncUtil.awaitConditionF(condition)
+      } yield succeed
+  }
+}

node-test/src/test/scala/org/bitcoins/node/SpvNodeWithWalletTest.scala

@@ -15,7 +15,7 @@ import org.scalatest.exceptions.TestFailedException
 import scala.concurrent.duration._
 import scala.concurrent.{Future, Promise}
 
-class NodeWithWalletTest extends NodeUnitTest {
+class SpvNodeWithWalletTest extends NodeUnitTest {
 
   /** Wallet config with data directory set to user temp directory */
   implicit override protected def config: BitcoinSAppConfig =

testkit/src/main/scala/org/bitcoins/testkit/node/NodeTestUtil.scala

@@ -121,18 +121,22 @@ abstract class NodeTestUtil extends P2PLogger {
       implicit ec: ExecutionContext): Future[Boolean] = {
     val rpcCountF = rpc.getBlockCount
     for {
-      count <- node.chainApiFromDb().flatMap(_.getFilterCount)
+      filterCount <- node.chainApiFromDb().flatMap(_.getFilterCount)
-      rpcCount <- rpcCountF
+      blockCount <- rpcCountF
-    } yield rpcCount == count
+    } yield {
+      blockCount == filterCount
+    }
   }
 
   def isSameBestFilterHeaderHeight(node: NeutrinoNode, rpc: BitcoindRpcClient)(
       implicit ec: ExecutionContext): Future[Boolean] = {
     val rpcCountF = rpc.getBlockCount
     for {
-      count <- node.chainApiFromDb().flatMap(_.getFilterHeaderCount)
+      filterHeaderCount <- node.chainApiFromDb().flatMap(_.getFilterHeaderCount)
-      rpcCount <- rpcCountF
+      blockCount <- rpcCountF
-    } yield rpcCount == count
+    } yield {
+      blockCount == filterHeaderCount
+    }
   }
 
   /** Checks if the given light client and bitcoind

testkit/src/main/scala/org/bitcoins/testkit/wallet/BitcoinSWalletTest.scala

@@ -83,6 +83,8 @@ trait BitcoinSWalletTest extends BitcoinSFixture with WalletLogger {
 
 object BitcoinSWalletTest extends WalletLogger {
 
+  lazy val initialFunds = 25.bitcoins
+
   case class WalletWithBitcoind(
       wallet: UnlockedWalletApi,
       bitcoind: BitcoindRpcClient)
@@ -173,7 +175,7 @@ object BitcoinSWalletTest extends WalletLogger {
     for {
       addr <- wallet.getNewAddress()
       tx <- bitcoind
-        .sendToAddress(addr, 25.bitcoins)
+        .sendToAddress(addr, initialFunds)
         .flatMap(bitcoind.getRawTransaction(_))
 
       _ <- bitcoind.getNewAddress.flatMap(bitcoind.generateToAddress(6, _))
@@ -181,7 +183,7 @@ object BitcoinSWalletTest extends WalletLogger {
       balance <- wallet.getBalance()
 
     } yield {
-      assert(balance >= 25.bitcoins)
+      assert(balance >= initialFunds)
       pair
     }
   }

wallet/src/main/scala/org/bitcoins/wallet/api/WalletApi.scala

@@ -6,7 +6,7 @@ import org.bitcoins.core.crypto._
 import org.bitcoins.core.currency.CurrencyUnit
 import org.bitcoins.core.hd.{AddressType, HDPurpose}
 import org.bitcoins.core.protocol.BitcoinAddress
-import org.bitcoins.core.protocol.blockchain.ChainParams
+import org.bitcoins.core.protocol.blockchain.{Block, ChainParams}
 import org.bitcoins.core.protocol.transaction.Transaction
 import org.bitcoins.core.wallet.fee.FeeUnit
 import org.bitcoins.wallet.HDUtil
@@ -52,6 +52,12 @@ trait LockedWalletApi extends WalletApi {
       transaction: Transaction,
       confirmations: Int): Future[LockedWalletApi]
 
+  /**
+    * Processes the give block, updating our DB state if it's relevant to us.
+    * @param block The block we're processing
+    */
+  def processBlock(block: Block, confirmations: Int): Future[LockedWalletApi]
+
   /** Gets the sum of all UTXOs in this wallet */
   def getBalance(): Future[CurrencyUnit] = {
     val confirmedF = getConfirmedBalance()

wallet/src/main/scala/org/bitcoins/wallet/internal/TransactionProcessing.scala

@@ -1,14 +1,14 @@
 package org.bitcoins.wallet.internal
 
-import org.bitcoins.core.protocol.transaction.Transaction
-import org.bitcoins.wallet._
-import org.bitcoins.wallet.models._
-import scala.concurrent.Future
-import org.bitcoins.core.protocol.transaction.TransactionOutput
-import org.bitcoins.wallet.api.AddUtxoSuccess
-import org.bitcoins.wallet.api.AddUtxoError
 import org.bitcoins.core.number.UInt32
+import org.bitcoins.core.protocol.blockchain.Block
+import org.bitcoins.core.protocol.transaction.{Transaction, TransactionOutput}
 import org.bitcoins.core.util.FutureUtil
+import org.bitcoins.wallet._
+import org.bitcoins.wallet.api.{AddUtxoError, AddUtxoSuccess}
+import org.bitcoins.wallet.models._
+
+import scala.concurrent.Future
 
 /** Provides functionality for processing transactions. This
   * includes importing UTXOs spent to our wallet, updating
@@ -35,6 +35,31 @@ private[wallet] trait TransactionProcessing extends KeyHandlingLogger {
 
   }
 
+  /** @inheritdoc */
+  override def processBlock(
+      block: Block,
+      confirmations: Int): Future[LockedWallet] = {
+    logger.info(
+      s"Processing block=${block.blockHeader.hash.flip} with confirmations=$confirmations")
+    val res = block.transactions.foldLeft(Future.successful(this)) {
+      (acc, transaction) =>
+        for {
+          _ <- acc
+          newWallet <- processTransaction(transaction, confirmations)
+        } yield {
+          newWallet
+        }
+    }
+    res.foreach(
+      _ =>
+        logger.info(
+          s"Finished processing of block=${block.blockHeader.hash.flip}."))
+    res.failed.foreach(e =>
+      logger.error(s"Error processing of block=${block.blockHeader.hash.flip}.",
+                   e))
+    res
+  }
+
   private[wallet] case class ProcessTxResult(
       updatedIncoming: List[SpendingInfoDb],
       updatedOutgoing: List[SpendingInfoDb])