github.com/mit-dci/lit@v0.0.0-20221102210550-8c3d3b49f2ce/test/testlib.py (about)

     1  import os
     2  import os.path as paths
     3  import time
     4  import signal
     5  import subprocess
     6  import logging
     7  import random
     8  import shutil
     9  
    10  import testutil
    11  import btcrpc
    12  import litrpc
    13  
    14  LIT_BIN = "%s/../lit" % paths.abspath(paths.dirname(__file__))
    15  
    16  REGTEST_COINTYPE = 257
    17  
    18  logger = logging.getLogger("testframework")
    19  
    20  next_unused_port = 11000
    21  def new_port():
    22      global next_unused_port
    23      port = next_unused_port
    24      next_unused_port += 1
    25      return port
    26  
    27  def get_root_data_dir():
    28      if 'LIT_ITEST_ROOT' in os.environ:
    29          return os.environ['LIT_ITEST_ROOT']
    30      else:
    31          return "%s/_data" % paths.abspath(paths.dirname(__file__))
    32  
    33  datadirnums = {}
    34  def new_data_dir(name):
    35      global datadirnums
    36      id = 0
    37      if name in datadirnums:
    38          id = datadirnums[name]
    39          datadirnums[name] += 1
    40      else:
    41          datadirnums[name] = 1 # set the next unused to "1"
    42      p = paths.join(get_root_data_dir(), name + str(id))
    43      os.makedirs(p, exist_ok=True)
    44      return p
    45  
    46  hexchars = "0123456789abcdef"
    47  
    48  # FIXME This doesn't work as expected anymore since IDs are global.
    49  next_id = 0
    50  def get_new_id():
    51      global next_id
    52      id = next_id
    53      next_id += 1
    54      return id
    55  
    56  class LitNode():
    57      def __init__(self, bcnode):
    58          self.bcnode = bcnode
    59          self.id = get_new_id()
    60          self.p2p_port = new_port()
    61          self.rpc_port = new_port()
    62          self.data_dir = new_data_dir("lit")
    63          self.peer_mapping = {}
    64          self.proc = None
    65  
    66          # Write a hexkey to the privkey file
    67          with open(paths.join(self.data_dir, "privkey.hex"), 'w+') as f:
    68              s = ''
    69              for _ in range(64):
    70                  s += hexchars[random.randint(0, len(hexchars) - 1)]
    71              print('Using key:', s)
    72              f.write(s + "\n")
    73  
    74          # Go and do the initial startup and sync.
    75          self.start()
    76  
    77      def start(self):
    78          # Sanity check.
    79          assert self.proc is None, "tried to start a node that is already started!"
    80  
    81          # See if we should print stdout
    82          outputredir = subprocess.DEVNULL
    83          ev_output_show = os.getenv("LIT_OUTPUT_SHOW", default="0")
    84          ev_show_id = os.getenv("LIT_ID_SHOW", default="X")
    85          if ev_output_show == "1" and (ev_show_id == "X" or ev_show_id == str(self.id)):
    86              outputredir = None
    87  
    88          # Now figure out the args to use and then start Lit.
    89          args = [
    90              LIT_BIN,
    91              "-vv",
    92              "--reg", "127.0.0.1:" + str(self.bcnode.p2p_port),
    93              "--tn3", "", # disable autoconnect
    94              "--dir", self.data_dir,
    95              "--unauthrpc",
    96              "--rpcport=" + str(self.rpc_port),
    97              "--noautolisten"
    98          ]
    99          penv = os.environ.copy()
   100          lkw = 'LIT_KEYFILE_WARN'
   101          if lkw not in penv:
   102              penv[lkw] = '0'
   103          self.proc = subprocess.Popen(args,
   104              stdin=subprocess.DEVNULL,
   105              stdout=outputredir,
   106              stderr=outputredir,
   107              env=penv)
   108  
   109          # Make the RPC client for future use, too.
   110          testutil.wait_until_port("localhost", self.rpc_port)
   111          self.rpc = litrpc.LitClient("localhost", str(self.rpc_port))
   112  
   113          # Make it listen to P2P connections!
   114          lres = self.rpc.Listen(Port=self.p2p_port)
   115          testutil.wait_until_port("localhost", self.p2p_port)
   116          self.lnid = lres["Adr"] # technically we do this more times than we have to, that's okay
   117  
   118      def get_sync_height(self):
   119          for bal in self.rpc.balance():
   120              if bal['CoinType'] == REGTEST_COINTYPE:
   121                  return bal['SyncHeight']
   122          return -1
   123  
   124      def connect_to_peer(self, other):
   125          addr = other.lnid + '@127.0.0.1:' + str(other.p2p_port)
   126          res = self.rpc.Connect(LNAddr=addr)
   127          self.update_peers()
   128          if 'PeerIdx' in res and self.peer_mapping[other.lnid] != res['PeerIdx']:
   129              raise AssertionError("new peer ID doesn't match reported ID (%s vs %s)" % (self.peer_mapping[other.lnid], res['PeerIdx']))
   130          other.update_peers()
   131  
   132      def get_peer_id(self, other):
   133          return self.peer_mapping[other.lnid]
   134  
   135      def make_new_addr(self):
   136          res = self.rpc.Address(NumToMake=1, CoinType=REGTEST_COINTYPE)
   137          return res['WitAddresses'][0]
   138  
   139      def update_peers(self):
   140          res = self.rpc.ListConnections()
   141          pm = {}
   142          for p in res['Connections']:
   143              pm[p['LitAdr']] = p['PeerNumber']
   144          self.peer_mapping = pm
   145  
   146      def get_balance_info(self, cointype=None):
   147          ct = REGTEST_COINTYPE
   148          if cointype is not None: # I had to do this because of reasons.
   149              ct = cointype
   150          for b in self.rpc.balance():
   151              if b['CoinType'] == ct:
   152                  return b
   153          return None
   154  
   155      def open_channel(self, peer, capacity, initialsend, cointype=None):
   156          ct = REGTEST_COINTYPE
   157          if cointype is not None: # I had to do thi because of reasons.
   158              ct = cointype
   159          res = self.rpc.FundChannel(
   160              Peer=self.get_peer_id(peer),
   161              CoinType=ct,
   162              Capacity=capacity,
   163              InitialSend=initialsend,
   164              Data=None) # maybe use [0 for _ in range(32)] or something?
   165          return res['ChanIdx']
   166  
   167      def resync(self):
   168          def ck_synced():
   169              return self.get_sync_height() == self.bcnode.get_block_height()
   170          testutil.wait_until(ck_synced, attempts=40, errmsg="node failing to resync!")
   171  
   172      def shutdown(self):
   173          if self.proc is not None:
   174              self.proc.kill()
   175              self.proc.wait()
   176              self.proc = None
   177          else:
   178              pass # do nothing I guess?
   179  
   180  class BitcoinNode():
   181      def __init__(self):
   182          self.p2p_port = new_port()
   183          self.rpc_port = new_port()
   184          self.data_dir = new_data_dir("bitcoind")
   185  
   186          # Actually start the bitcoind
   187          args = [
   188              "bitcoind",
   189              "-regtest",
   190              "-server",
   191              "-printtoconsole",
   192              "-datadir=" + self.data_dir,
   193              "-port=" + str(self.p2p_port),
   194              "-rpcuser=rpcuser",
   195              "-rpcpassword=rpcpass",
   196              "-rpcport=" + str(self.rpc_port),
   197          ]
   198          self.proc = subprocess.Popen(args,
   199              stdout=subprocess.DEVNULL,
   200              stderr=subprocess.DEVNULL)
   201  
   202          # Make the RPC client for it.
   203          testutil.wait_until_port("localhost", self.rpc_port)
   204          testutil.wait_until_port("localhost", self.p2p_port)
   205          self.rpc = btcrpc.BtcClient("localhost", self.rpc_port, "rpcuser", "rpcpass")
   206  
   207          # Make sure that we're actually ready to accept RPC calls.
   208          def ck_ready():
   209              bci = self.rpc.getblockchaininfo() # just need "some call" that'll fail if we're not ready
   210              if 'code' in bci and bci['code'] <= 0:
   211                  return False
   212              else:
   213                  return True
   214          testutil.wait_until(ck_ready, errmsg="took too long to load wallet")
   215  
   216          # Activate SegWit (apparently this is how you do it)
   217          self.rpc.generate(500)
   218          def ck_segwit():
   219              bci = self.rpc.getblockchaininfo()
   220              try:
   221                  return bci["bip9_softforks"]["segwit"]["status"] == "active"
   222              except:
   223                  return False
   224          testutil.wait_until(ck_segwit, errmsg="couldn't activate segwit")
   225  
   226      def get_block_height(self):
   227          return self.rpc.getblockchaininfo()['blocks']
   228  
   229      def shutdown(self):
   230          if self.proc is not None:
   231              self.proc.kill()
   232              self.proc.wait()
   233              self.proc = None
   234          else:
   235              pass # do nothing I guess?
   236  
   237  class TestEnv():
   238      def __init__(self, litcnt):
   239          logger.info("starting nodes...")
   240          self.bitcoind = BitcoinNode()
   241          self.lits = []
   242          for i in range(litcnt):
   243              node = LitNode(self.bitcoind)
   244              self.lits.append(node)
   245          logger.info("started nodes!  syncing...")
   246  
   247          time.sleep(0.1)
   248  
   249          # Sync the nodes
   250          try:
   251              self.generate_block(count=0)
   252          except Exception as e:
   253              logger.warning("probem syncing nodes, exiting (" + str(e) + ")")
   254              self.shutdown()
   255          logger.info("nodes synced!")
   256  
   257      def new_lit_node(self):
   258          node = LitNode(self.bitcoind)
   259          self.lits.append(node)
   260          self.generate_block(count=0) # Force it to wait for sync.
   261          return node
   262  
   263      def generate_block(self, count=1):
   264          if count > 0:
   265              self.bitcoind.rpc.generate(count)
   266          h = self.bitcoind.get_block_height()
   267          def ck_lits_synced():
   268              for l in self.lits:
   269                  sh = l.get_sync_height()
   270                  if sh != h:
   271                      return False
   272              return True
   273          testutil.wait_until(ck_lits_synced, errmsg="lits aren't syncing to bitcoind")
   274  
   275      def get_height(self):
   276          return self.bitcoind.get_block_height()
   277  
   278      def shutdown(self):
   279          for l in self.lits:
   280              l.shutdown()
   281          self.bitcoind.shutdown()
   282  
   283  def clean_data_dir():
   284      datadir = get_root_data_dir()
   285      shutil.rmtree(datadir)