From 5afd446bb3a9cf1d76efe4c5c2d83ed3a285ef82 Mon Sep 17 00:00:00 2001
From: Lex Berezhny <lex@damoti.com>
Date: Thu, 16 Aug 2018 12:40:54 -0400
Subject: [PATCH] fixing wallet unit tests test_ledger.py and
 test_transactions.py

---
 tests/unit/wallet/test_ledger.py      | 36 ++++++---------------------
 tests/unit/wallet/test_transaction.py | 26 ++++++++++++++-----
 2 files changed, 28 insertions(+), 34 deletions(-)

diff --git a/tests/unit/wallet/test_ledger.py b/tests/unit/wallet/test_ledger.py
index 7d5bf79ce..05a47bf0c 100644
--- a/tests/unit/wallet/test_ledger.py
+++ b/tests/unit/wallet/test_ledger.py
@@ -1,45 +1,25 @@
 from twisted.internet import defer
 from twisted.trial import unittest
-from lbrynet import conf
 from lbrynet.wallet.account import Account
-from lbrynet.wallet.database import WalletDatabase
 from lbrynet.wallet.transaction import Transaction, Output, Input
 from lbrynet.wallet.ledger import MainNetLedger
 from torba.wallet import Wallet
 
 
-class MockHeaders:
-    def __init__(self, ledger):
-        self.ledger = ledger
-        self.height = 1
-
-    def __len__(self):
-        return self.height
-
-    def __getitem__(self, height):
-        return {'merkle_root': 'abcd04'}
-
-
-class MainNetTestLedger(MainNetLedger):
-    headers_class = MockHeaders
-    network_name = 'unittest'
-
-    def __init__(self):
-        super(MainNetLedger, self).__init__({
-            'db': WalletDatabase(':memory:')
-        })
-
-
 class LedgerTestCase(unittest.TestCase):
 
     def setUp(self):
-        conf.initialize_settings(False)
-        self.ledger = MainNetTestLedger()
+        super().setUp()
+        self.ledger = MainNetLedger({
+            'db': MainNetLedger.database_class(':memory:'),
+            'headers': MainNetLedger.headers_class(':memory:')
+        })
         self.account = Account.generate(self.ledger, Wallet(), "lbryum")
-        return self.ledger.db.start()
+        return self.ledger.db.open()
 
     def tearDown(self):
-        return self.ledger.db.stop()
+        super().tearDown()
+        return self.ledger.db.close()
 
 
 class BasicAccountingTests(LedgerTestCase):
diff --git a/tests/unit/wallet/test_transaction.py b/tests/unit/wallet/test_transaction.py
index e988bf3f8..399e73317 100644
--- a/tests/unit/wallet/test_transaction.py
+++ b/tests/unit/wallet/test_transaction.py
@@ -5,11 +5,8 @@ from twisted.internet import defer
 from torba.constants import CENT, COIN, NULL_HASH32
 from torba.wallet import Wallet
 
-from lbrynet.wallet.account import Account
 from lbrynet.wallet.ledger import MainNetLedger
-from lbrynet.wallet.database import WalletDatabase
 from lbrynet.wallet.transaction import Transaction, Output, Input
-from lbrynet.wallet.manager import LbryWalletManager
 
 
 FEE_PER_BYTE = 50
@@ -41,7 +38,16 @@ def get_claim_transaction(claim_name, claim=b''):
 class TestSizeAndFeeEstimation(unittest.TestCase):
 
     def setUp(self):
-        self.ledger = MainNetLedger({'db': WalletDatabase(':memory:')})
+        super().setUp()
+        self.ledger = MainNetLedger({
+            'db': MainNetLedger.database_class(':memory:'),
+            'headers': MainNetLedger.headers_class(':memory:')
+        })
+        return self.ledger.db.open()
+
+    def tearDown(self):
+        super().tearDown()
+        return self.ledger.db.close()
 
     def test_output_size_and_fee(self):
         txo = get_output()
@@ -216,8 +222,16 @@ class TestTransactionSerialization(unittest.TestCase):
 class TestTransactionSigning(unittest.TestCase):
 
     def setUp(self):
-        self.ledger = MainNetLedger({'db': WalletDatabase(':memory:')})
-        return self.ledger.db.start()
+        super().setUp()
+        self.ledger = MainNetLedger({
+            'db': MainNetLedger.database_class(':memory:'),
+            'headers': MainNetLedger.headers_class(':memory:')
+        })
+        return self.ledger.db.open()
+
+    def tearDown(self):
+        super().tearDown()
+        return self.ledger.db.close()
 
     @defer.inlineCallbacks
     def test_sign(self):