source: branches/numpy_misc/tools/acceptance_tests/test_utils.py @ 7026

#!/bin/env python

'''Utility functions for the cluster acceptance test suite.'''

import os
import sys
import time


# information about each cluster.
# 'node_stem' key defines a compute node name.
# update this if adding a new cluster.
# should get this from the machines_<cluster> file?
Cluster_Info = {'cyclone': {'num_nodes': 23,
                            'node_stem': 'compute-1-%d',
                            'bad_nodes': [13, 22]},
                # 13 bad - 'no route to host'
                # 22 bad - forgets password, can't find file
                'tornado': {'num_nodes': 20,
                            'node_stem': 'compute-0-%d',
                            'bad_nodes': []}
               }


def log(logfile, msg=''):
    '''Write a message to the log file.'''

    fd = open(logfile, 'a')
    fd.write(msg)
    fd.close()


def log_nl(logfile, msg=''):
    '''Write a message to the log file, with a newline added.'''

    log(logfile, msg)
    log(logfile, '\n')


def log_print(logfile, msg=''):
    '''Write a message to the log file and stdout.

    The message doesn't have a '\n' added to it.
    '''

    log(logfile, msg)
    print msg,
    sys.stdout.flush()


def log_print_nl(logfile, msg=''):
    '''Write a message to the log file and stdout, with a newline added.

    Appends '\n' to msg.
    '''

    log_nl(logfile, msg)
    print msg
    sys.stdout.flush()


def header(logfile, module, msg):
    '''Write testlet header to log and stdout.'''

    # current time
    time_now = time.strftime('%H:%M:%S %Y-%m-%d')

    log_print_nl(logfile, '#' * 80)
    log_print_nl(logfile, '# ' + time_now + ' - ' + msg)
    log_print_nl(logfile, '# Module: %s' % module)
    log_print_nl(logfile, '#' * 80)
    log_print_nl(logfile)


def footer(logfile, start_time=None):
    '''Write testlet footer to log and stdout.'''

    if start_time:
        end_time = time.time()
        msg = 'Test took %.1f s' % (end_time - start_time)
    else:
        msg = 'Finished at %s' % time.ctime()
        
    log_nl(logfile)
    log_nl(logfile, '# ' + msg)
    log_print_nl(logfile)


def get_hostname():
    '''Return (<host>, <domain>) for machine.

    For example, on the GA cyclone main node, return ('cyclone', 'agso.gov.au').
    '''

    fd = os.popen('hostname')
    hostname = fd.read()
    fd.close()

    try:
        result = hostname.strip().split('.', 1)
        if len(result) < 2:
            result.append('')
    except:
        result = (hostname, '')

    return result


def get_cluster_info(cluster):
    '''Get dictionary containing information for 'cluster'.

    If 'cluster' is not recognized, return None.
    '''

    return Cluster_Info.get(cluster, None)


def get_node_numbers(cluster, strip_bad_nodes=True):
    '''Get a list of node numbers for the cluster.

    If 'strip_bad_nodes' is True remove bad nodes from the result.
    '''

    result = []
    info = Cluster_Info[cluster]
    for n in xrange(info['num_nodes']):
        if strip_bad_nodes:
            if n in info['bad_nodes']:
                continue
        result.append(n)

    return result