BalsamTestCase.py 2.52 KB
Newer Older
1 2
import os
import unittest
Michael Salim's avatar
Michael Salim committed
3
import subprocess
4
import time
5 6 7 8 9

from django.core.management import call_command
from django import db
from django.conf import settings

Michael Salim's avatar
Michael Salim committed
10 11
from balsam.models import BalsamJob, ApplicationDefinition

12 13 14
class BalsamTestCase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
Michael Salim's avatar
Michael Salim committed
15
        assert db.connection.settings_dict['NAME'].endswith('test_db.sqlite3')
16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
        call_command('makemigrations',interactive=False,verbosity=0)
        call_command('migrate',interactive=False,verbosity=0)
        assert os.path.exists(settings.DATABASES['default']['NAME'])

    @classmethod
    def tearDownClass(cls):
        pass

    def setUp(self):
        pass # to be implemented by test cases

    def tearDown(self):
        if not db.connection.settings_dict['NAME'].endswith('test_db.sqlite3'):
            raise RuntimeError("Test DB not configured")
        call_command('flush',interactive=False,verbosity=0)
Michael Salim's avatar
Michael Salim committed
31 32 33 34 35 36

def cmdline(cmd,envs=None,shell=True):
    '''Return string output from a command line'''
    p = subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE,
                         stderr=subprocess.STDOUT,env=envs)
    return p.communicate()[0].decode('utf-8')
37 38 39 40 41 42 43 44

def poll_until_returns_true(function, *, args=(), period=1.0, timeout=12.0):
    start = time.time()
    while time.time() - start < timeout:
        result = function(*args)
        if result: break
        else: time.sleep(period)
    return result
Michael Salim's avatar
Michael Salim committed
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81

def create_job(*, name='', app='', direct_command='', site=settings.BALSAM_SITE, num_nodes=1,
               ranks_per_node=1, args='', workflow='', envs={}, state='CREATED'):

    if app and direct_command:
        raise ValueError("Cannot have both application and direct command")

    job = BalsamJob()
    job.name = name
    job.application = app
    job.direct_command = direct_command
    
    job.allowed_work_sites = site
    
    job.num_nodes = num_nodes
    job.ranks_per_node = ranks_per_node
    job.application_args = args
    
    job.workflow = workflow
    job.environ_vars = ':'.join(f'{k}={v}' for k,v in envs.items())
    job.state = state
    job.save()
    job.create_working_path()
    return job

def create_app(*, name='', description='', executable='', preproc='',
               postproc='', envs={}):

    app = ApplicationDefinition()
    app.name = name
    app.description = description
    app.executable = executable
    app.default_preprocess = preproc
    app.default_postprocess = postproc
    app.environ_vars = ':'.join(f'{k}={v}' for k,v in envs.items())
    app.save()
    return app