def attempt(message, exception_types=Exception, include_traceback=False):
"""Context manager to be used to wrap statements in CLI that can throw exceptions.
:param message: the message to print before yielding
:param include_traceback: boolean, if True, will also print traceback if an exception is caught
"""
import sys
import traceback
echo.echo_info(message, nl=False)
try:
yield
except exception_types as exception: # pylint: disable=broad-except
echo.echo_highlight(' [FAILED]', color='error', bold=True)
message = str(exception)
if include_traceback:
message += f"\n{''.join(traceback.format_exception(*sys.exc_info()))}"
echo.echo_critical(message)
else:
echo.echo_highlight(' [OK]', color='success', bold=True)
def computer_test(user, print_traceback, computer):
"""
Test the connection to a computer.
It tries to connect, to get the list of calculations on the queue and
to perform other tests.
"""
import traceback
from aiida import orm
from aiida.common.exceptions import NotExistent
# Set a user automatically if one is not specified in the command line
if user is None:
user = orm.User.objects.get_default()
echo.echo_info(
f'Testing computer<{computer.label}> for user<{user.email}>...')
try:
authinfo = computer.get_authinfo(user)
except NotExistent:
echo.echo_critical(
f'Computer<{computer.label}> is not yet configured for user<{user.email}>'
)
if not authinfo.enabled:
echo.echo_warning(
f'Computer<{computer.label}> is disabled for user<{user.email}>')
click.confirm('Do you really want to test it?', abort=True)
scheduler = authinfo.computer.get_scheduler()
transport = authinfo.get_transport()
# STARTING TESTS HERE
num_failures = 0
num_tests = 0
tests = {
_computer_test_no_unexpected_output: 'Checking for spurious output',
_computer_test_get_jobs: 'Getting number of jobs from scheduler',
_computer_get_remote_username: '******',
_computer_create_temp_file: 'Creating and deleting temporary file'
}
try:
echo.echo('* Opening connection... ', nl=False)
with transport:
num_tests += 1
echo.echo_highlight('[OK]', color='success')
scheduler.set_transport(transport)
for test, test_label in tests.items():
echo.echo(f'* {test_label}... ', nl=False)
num_tests += 1
try:
success, message = test(transport=transport,
scheduler=scheduler,
authinfo=authinfo)
except Exception as exception: # pylint:disable=broad-except
success = False
message = f'{exception.__class__.__name__}: {str(exception)}'
if print_traceback:
message += '\n Full traceback:\n'
message += '\n'.join([
' {}'.format(l)
for l in traceback.format_exc().splitlines()
])
else:
message += '\n Use the `--print-traceback` option to see the full traceback.'
if not success:
num_failures += 1
if message:
echo.echo_highlight('[Failed]: ',
color='error',
nl=False)
echo.echo(message)
else:
echo.echo_highlight('[Failed]', color='error')
else:
if message:
echo.echo_highlight('[OK]: ',
color='success',
nl=False)
echo.echo(message)
else:
echo.echo_highlight('[OK]', color='success')
if num_failures:
echo.echo_warning(
f'{num_failures} out of {num_tests} tests failed')
else:
echo.echo_success(f'all {num_tests} tests succeeded')
except Exception as exception: # pylint:disable=broad-except
echo.echo_highlight('[FAILED]: ', color='error', nl=False)
message = 'Error while trying to connect to the computer'
if print_traceback:
message += '\n Full traceback:\n'
message += '\n'.join([
' {}'.format(l) for l in traceback.format_exc().splitlines()
])
else:
message += '\n Use the `--print-traceback` option to see the full traceback.'
echo.echo(message)
echo.echo_warning(f'{1} out of {num_tests} tests failed')