stacks-puppet-node/integration_tests/bin/blockstack-test-scenario

#!/usr/bin/env python 

import os
import sys 
import subprocess
import signal
import shutil
import time
import atexit
import errno
import importlib
import traceback

blockstack = None
blockstackd = None
mock_bitcoind = None
testlib = None

# enable all tests
os.environ['BLOCKSTACK_TEST'] = '1'
os.environ['BLOCKSTACK_CLIENT_TEST_ALTERNATIVE_CONFIG'] = '1'

import pybitcoin

import virtualchain 
log = virtualchain.get_logger("blockstack-test-scenario")

mock_bitcoind_connection = None
TEST_RPC_PORT = 16264
TEST_CLIENT_RPC_PORT = 16268
BLOCKSTACK_STORAGE_DRIVERS = "disk"

if os.environ.get("BLOCKSTACK_STORAGE_DRIVERS", None) is not None:
    BLOCKSTACK_STORAGE_DRIVERS = os.environ.get("BLOCKSTACK_STORAGE_DRIVERS")

DEFAULT_SERVER_INI_TEMPLATE = """
[bitcoind]
passwd = blockstacksystem
server = bitcoin.blockstack.org
port = 8332
use_https = True
user = blockstack
mock = True
save_file = @MOCK_SAVE_FILE@

[blockstack]
server_version = 0.0.13.0
rpc_port = %s
backup_frequency = 3
backup_max_age = 30
blockchain_proxy = True
serve_zonefiles = True
serve_profiles = True
zonefiles = @ZONEFILES@
analytics_key = abcdef0123456789
zonefile_storage_drivers = disk
profile_storage_drivers = disk
atlas = False
""" % TEST_RPC_PORT

DEFAULT_CLIENT_INI_TEMPLATE = """
[blockstack-client]
client_version = 0.0.13.0
server = localhost
port = %s
metadata = @CLIENT_METADATA@
storage_drivers = @CLIENT_STORAGE_DRIVERS@
storage_drivers_required_write = disk,blockstack-server
blockchain_headers = @CLIENT_BLOCKCHAIN_HEADERS@
advanced_mode = true
api_endpoint_port = %s
rpc_token = a653b93e696a998f85f8fd2b241ff4dfcb5dd978fe1da26c413a4c2abf90321b
poll_interval = 1
queue_path = @CLIENT_QUEUE_PATH@
rpc_detach = True
blockchain_reader = mock_utxo
blockchain_writer = mock_utxo
anonymous_statistics = False

[blockchain-reader]
utxo_provider = mock_utxo
initial_utxos = @MOCK_INITIAL_UTXOS@
save_file = @MOCK_SAVE_FILE@

[blockchain-writer]
utxo_provider = mock_utxo
initial_utxos = @MOCK_INITIAL_UTXOS@
save_file = @MOCK_SAVE_FILE

[bitcoind]
passwd = blockstacksystem
server = bitcoin.blockstack.org
port = 8332
use_https = True
user = blockstack
mock = True
save_file = @MOCK_SAVE_FILE@
blockchain_headers = @CLIENT_BLOCKCHAIN_HEADERS@
""" % (TEST_RPC_PORT, TEST_CLIENT_RPC_PORT)


def load_scenario( scenario_name ):
    """
    Load up a scenario, and validate it.
    A scenario is a python file with:
    * a global variable 'wallet' that is a dict
    which maps private keys to their initial values.
    * a global variable 'consensus' that represents 
    the initial consensus hash.
    * a callable called 'scenario' that takes the 
    wallet as an argument and runs the test.
    * a callable called 'check' that takes the state 
    engine as an argument and checks it for correctness.
    """

    # strip .py from scenario name 
    if scenario_name.endswith(".py"):
        scenario_name = scenario_name[:-3]

    try:
        scenario = importlib.import_module( scenario_name )
    except ImportError, ie:
        raise Exception("Failed to import '%s'." % scenario_name )

    # validate 
    if not hasattr( scenario, "wallets" ):
        # default empty wallet 
        log.warning("Empty wallet for scenario '%s'" % scenario_name )
        scenario.wallets = {}

    if not hasattr( scenario, "consensus" ):
        # default consensus hash 
        log.warning("No consensus hash for '%s'" % scenario_name )
        scenario.consensus = "00" * 16

    if not hasattr( scenario, "scenario" ):
        # not a valid test 
        log.error("Invalid scenario '%s': no 'scenario' method" % scenario_name )
        return None 

    if not hasattr( scenario, "check" ):
        # not a valid test 
        log.error("Invalid scenario '%s': no 'check' method" % scenario_name )
        return None 

    return scenario
   

def generate_config_file( scenario, path, template, extra_fields):
    """
    Generate the config file to use with this test scenario.
    Write it to path.
    """

    client_metadata = extra_fields.get("CLIENT_METADATA", "")
    queue_path = extra_fields.get("CLIENT_QUEUE_PATH", "")
    zonefiles_dir = extra_fields.get("ZONEFILES", None)
    storage_drivers = extra_fields.get("CLIENT_STORAGE_DRIVERS", BLOCKSTACK_STORAGE_DRIVERS)
    if zonefiles_dir is None:
        zonefiles_dir = "/tmp/blockstack-test-zonefiles"

    config_txt = None
    if template is not None:
        config_txt = template[:]

    else:
        raise Exception("No config template")

    initial_utxo_str = ",".join( ["%s:%s" % (w.privkey, w.value) for w in scenario.wallets] )
    mock_bitcoind_save_path = "/tmp/mock_bitcoind.dat"
    mock_bitcoind_headers_path = "/tmp/mock_bitcoind.dat.spvheaders"

    config_txt = config_txt.replace( "@MOCK_INITIAL_UTXOS@", initial_utxo_str )
    config_txt = config_txt.replace( "@MOCK_SAVE_FILE@", mock_bitcoind_save_path )
    config_txt = config_txt.replace( "@CLIENT_BLOCKCHAIN_HEADERS@", mock_bitcoind_headers_path )
    config_txt = config_txt.replace( "@CLIENT_METADATA@", client_metadata ) 
    config_txt = config_txt.replace( "@CLIENT_QUEUE_PATH@", queue_path )
    config_txt = config_txt.replace( "@ZONEFILES@", zonefiles_dir )
    config_txt = config_txt.replace( "@CLIENT_STORAGE_DRIVERS@", storage_drivers )

    with open( path, "w" ) as f:
        f.write( config_txt )
        f.flush()

    return 0


def network_start():
    """
    Start RPC services
    """
    blockstackd.rpc_start(TEST_RPC_PORT)


def network_stop():
    """
    Stop RPC services
    """
    blockstackd.rpc_stop()


def atexit_cleanup( mock_bitcoind_save_path, mock_bitcoind_headers_path, client_config_path ):
    """
    Clean up a scenario
    """
    try:
        os.unlink(mock_bitcoind_save_path)
    except:
        pass

    try:
        os.unlink(mock_bitcoind_headers_path)
    except:
        pass

    if client_config_path is not None:
        client_config_dir = os.path.dirname(client_config_path)
        blockstack_client.rpc.local_rpc_stop(client_config_dir)


def run_scenario( scenario, config_file, client_config_file, interactive=False, blocktime=10 ):
    """
    Run a test scenario:
    * set up the virtualchain to use our mock UTXO provider and mock bitcoin blockchain
    * seed it with the initial values in the wallet 
    * set the initial consensus hash 
    * start the API server
    * run the scenario method
    * run the check method
    """

    mock_bitcoind_save_path = "/tmp/mock_bitcoind.dat"
    spv_headers_path = mock_bitcoind_save_path + ".spvheaders"

    atexit_cleanup( mock_bitcoind_save_path, spv_headers_path, None )
    atexit.register( atexit_cleanup, mock_bitcoind_save_path, spv_headers_path, client_config_file )

    # use mock bitcoind
    worker_env = mock_bitcoind.make_worker_env( mock_bitcoind, mock_bitcoind_save_path )
    worker_env['BLOCKSTACK_TEST'] = "1"

    print worker_env

    # tell our subprocesses that we're testing 
    os.environ.update( worker_env )

    if os.environ.get("PYTHONPATH", None) is not None:
        worker_env["PYTHONPATH"] = os.environ["PYTHONPATH"]

    # virtualchain defaults...
    virtualchain.setup_virtualchain( impl=blockstack_state_engine, bitcoind_connection_factory=mock_bitcoind.connect_mock_bitcoind, index_worker_env=worker_env )

    # set up blockstack
    # NOTE: utxo_opts encodes the mock-bitcoind options
    server_opts = blockstack.lib.config.configure(config_file=config_file, interactive=False)
    blockstack_opts = server_opts['blockstack']
    bitcoin_opts = server_opts['bitcoind']

    client_opts = blockstack_client.config.configure(config_file=client_config_file, interactive=False)
    utxo_opts = client_opts['blockchain-reader']
   
    # override multiprocessing options to ensure single-process behavior 
    utxo_opts['multiprocessing_num_procs'] = 1 
    utxo_opts['multiprocessing_num_blocks'] = 10

    # pass along extra arguments
    utxo_opts['save_file'] = mock_bitcoind_save_path
    utxo_opts['blockchain'] = 'blockstack_integration_tests.mock_bitcoind'
    utxo_opts['blockchain_server'] = '127.0.0.1'
    utxo_opts['blockchain_port'] = 8332

    print ""
    print "blockstack opts"
    print json.dumps( blockstack_opts, indent=4 )

    print ""
    print "blockchain opts"
    print json.dumps( bitcoin_opts, indent=4 )

    print ""
    print "blockchain service opts"
    print json.dumps( utxo_opts, indent=4 )

    # save headers as well 
    utxo_opts['spv_headers_path'] = spv_headers_path
    with open( utxo_opts['spv_headers_path'], "w" ) as f:
        # write out "initial" headers, up to the first block
        empty_header = ("00" * 81).decode('hex')
        for i in xrange(0, blockstack.FIRST_BLOCK_MAINNET ): 
            f.write( empty_header )

    print "mock blockchain chainstate save path: %s" % mock_bitcoind_save_path
    print "mock blockchain headers: %s" % spv_headers_path

    blockstackd.set_blockstack_opts( blockstack_opts )
    blockstackd.set_bitcoin_opts( bitcoin_opts )

    db = blockstackd.get_db_state(disposition=blockstackd.DISPOSITION_RW)
    bitcoind = mock_bitcoind.connect_mock_bitcoind( utxo_opts )

    sync_virtualchain_upcall = lambda: virtualchain.sync_virtualchain( utxo_opts, bitcoind.getblockcount(), db )
    mock_utxo = mock_utxo_provider.MockUTXOProvider( bitcoind )
    working_dir = virtualchain.get_working_dir()

    print "working_dir: %s" % working_dir
    print "UTXO client: %s" % mock_utxo

    # set up test environment
    testlib.set_utxo_opts( utxo_opts )
    testlib.set_utxo_client( mock_utxo )
    testlib.set_bitcoind( bitcoind )
    testlib.set_state_engine( db )

    test_env = {
        "sync_virtualchain_upcall": sync_virtualchain_upcall,
        "working_dir": working_dir,
        "bitcoind": bitcoind,
        "bitcoind_save_path": mock_bitcoind_save_path
    }

    # start taking RPC requests
    network_start()

    # sync initial utxos 
    testlib.next_block( **test_env )


    # load the scenario into the mock blockchain and mock utxo provider
    try:
        scenario.scenario( scenario.wallets, **test_env )

    except Exception, e:
        log.exception(e)
        traceback.print_exc()
        log.error("Failed to run scenario '%s'" % scenario.__name__)
        network_stop()
        return False

    # run the checks on the database
    try:
        rc = scenario.check( db )
    except Exception, e:
        log.exception(e)
        traceback.print_exc()
        log.error("Failed to run tests '%s'" % scenario.__name__)
        network_stop()
        return False 
    
    if not rc:
        network_stop()
        return rc

    # do any more interactive tests
    if interactive:
        log.info("Keeping test server online for testing purposes.")
        log.info("Blocktime is %s second(s)" % blocktime)
        while True:
            try:
                db = blockstackd.get_db_state(disposition=blockstackd.DISPOSITION_RW)
                testlib.set_state_engine( db )
                time.sleep(blocktime)
                testlib.next_block( **test_env )
            except KeyboardInterrupt:
                log.info('Resume tests')
                break

    log.info("Scenario checks passed; verifying history")

    # run database integrity check at each block 
    rc = testlib.check_history( db )
    if not rc:
        network_stop()
        return rc

    log.info("History check passes!")

    # run snv at each name 
    rc = testlib.snv_all_names( db )
    if not rc:
        network_stop()
        return rc

    log.info("SNV check passes!")
    network_stop()
    TEST_RESULT = True
    return True


if __name__ == "__main__":

    if len(sys.argv) < 2:
        print >> sys.stderr, "Usage: %s [--interactive [blocktime]] [scenario.import.path] [OPTIONAL: working dir]"
        sys.exit(1)

    working_dir = None
    interactive = False
    blocktime = 10

    if len(sys.argv) > 2:
        if sys.argv[1] == '--interactive':
            sys.argv.pop(1)
            try:
                blocktime = int(sys.argv[1])
                sys.argv.pop(1)
            except:
                pass
        
            interactive = True
            log.debug("Interactive session; blocktime is %s" % blocktime)

    if len(sys.argv) > 2:
        working_dir = sys.argv[2]

    else:
        working_dir = "/tmp/blockstack-run-scenario.%s" % sys.argv[1]
        if os.path.exists(working_dir):
            shutil.rmtree(working_dir)
   
    if not os.path.exists(working_dir):
        os.makedirs(working_dir)

    client_working_dir = os.path.join(working_dir, "client")
    client_metadata = os.path.join(client_working_dir, "metadata")
    client_queue_path = os.path.join( client_working_dir, "queues.db" )
    config_file = os.path.join( working_dir, "blockstack-server.ini" ) 
    client_config_file = os.path.join( client_working_dir, "client.ini" )

    os.makedirs(client_working_dir)
    os.makedirs(client_metadata) 

    # export to test 
    os.environ["BLOCKSTACK_CLIENT_CONFIG"] = client_config_file
    os.environ["BLOCKSTACK_SERVER_CONFIG"] = config_file
    os.environ['VIRTUALCHAIN_WORKING_DIR'] = working_dir

    # *now* we can import blockstack 
    import blockstack
    import blockstack.blockstackd as blockstackd
    from blockstack.lib import *

    from blockstack.lib import nameset as blockstack_state_engine

    import blockstack_integration_tests.mock_bitcoind as mock_bitcoind
    import blockstack_integration_tests.mock_utxo_provider as mock_utxo_provider
    import blockstack_integration_tests.scenarios.testlib as testlib

    # load up the scenario 
    scenario = load_scenario( sys.argv[1] )
    if scenario is None:
        print "Failed to load '%s'" % sys.argv[1]
        sys.exit(1)

    # did the scenario change this variable?
    storage_drivers = BLOCKSTACK_STORAGE_DRIVERS
    if os.environ.get("CLIENT_STORAGE_DRIVERS", None) is not None:
        storage_drivers = os.environ.get("CLIENT_STORAGE_DRIVERS")

    # generate config file
    rc = generate_config_file( scenario, config_file, DEFAULT_SERVER_INI_TEMPLATE, {"ZONEFILES": os.path.join(working_dir, "zonefiles")} )
    if rc != 0:
        log.error("failed to write config file: exit %s" % rc)
        sys.exit(1)

    # generate config file for the client
    rc = generate_config_file( scenario, client_config_file, DEFAULT_CLIENT_INI_TEMPLATE, {"CLIENT_METADATA": client_metadata, "CLIENT_QUEUE_PATH": client_queue_path, "CLIENT_STORAGE_DRIVERS": storage_drivers})
    if rc != 0:
        log.error("failed to write config file: exit %s" % rc)
        sys.exit(1)

    # run the test 
    rc = run_scenario( scenario, config_file, client_config_file, interactive=interactive, blocktime=blocktime )
   
    if rc:
        print "SUCCESS %s" % scenario.__name__
        # shutil.rmtree( working_dir )
        sys.exit(0)
    else:
        print >> sys.stderr, "FAILURE %s" % scenario.__name__
        print >> sys.stderr, "Test output in %s" % working_dir
        sys.exit(1)