<s> import subprocess <EOL> import os <EOL> import os . path <EOL> from glob import glob <EOL> import unittest <EOL> import logging <EOL> import luigi <EOL> from luigi . contrib . sge import SGEJobTask , _parse_qstat_state , _clean_task_id <EOL> DEFAULT_HOME = '<STR_LIT>' <EOL> logger = logging . getLogger ( '<STR_LIT>' ) <EOL> QSTAT_OUTPUT = """<STR_LIT>""" <EOL> def on_sge_master ( ) : <EOL> try : <EOL> subprocess . check_output ( '<STR_LIT>' , shell = True ) <EOL> return True <EOL> except subprocess . CalledProcessError : <EOL> return False <EOL> class TestSGEWrappers ( unittest . TestCase ) : <EOL> def test_track_job ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . assertEqual ( _parse_qstat_state ( QSTAT_OUTPUT , <NUM_LIT:1> ) , '<STR_LIT:r>' ) <EOL> self . assertEqual ( _parse_qstat_state ( QSTAT_OUTPUT , <NUM_LIT:2> ) , '<STR_LIT>' ) <EOL> self . assertEqual ( _parse_qstat_state ( QSTAT_OUTPUT , <NUM_LIT:3> ) , '<STR_LIT:t>' ) <EOL> self . assertEqual ( _parse_qstat_state ( '<STR_LIT>' , <NUM_LIT:1> ) , '<STR_LIT:u>' ) <EOL> self . assertEqual ( _parse_qstat_state ( '<STR_LIT>' , <NUM_LIT:4> ) , '<STR_LIT:u>' ) <EOL> def test_clean_task_id ( self ) : <EOL> task_id = '<STR_LIT>' <EOL> cleaned_id = '<STR_LIT>' <EOL> self . assertEqual ( _clean_task_id ( task_id ) , cleaned_id ) <EOL> class TestJobTask ( SGEJobTask ) : <EOL> '''<STR_LIT>''' <EOL> i = luigi . Parameter ( ) <EOL> def work ( self ) : <EOL> logger . info ( '<STR_LIT>' ) <EOL> with open ( self . output ( ) . path , '<STR_LIT:w>' ) as f : <EOL> f . write ( '<STR_LIT>' ) <EOL> def output ( self ) : <EOL> return luigi . LocalTarget ( os . path . join ( DEFAULT_HOME , '<STR_LIT>' + str ( self . i ) ) ) <EOL> class TestSGEJob ( unittest . TestCase ) : <EOL> '''<STR_LIT>''' <EOL> def test_run_job ( self ) : <EOL> if on_sge_master ( ) : <EOL> outfile = os . path . join ( DEFAULT_HOME , '<STR_LIT>' ) <EOL> tasks = [ TestJobTask ( i = str ( i ) , n_cpu = <NUM_LIT:1> ) for i in range ( <NUM_LIT:3> ) ] <EOL> luigi . build ( tasks , local_scheduler = True , workers = <NUM_LIT:3> ) <EOL> self . assertTrue ( os . path . exists ( outfile ) ) <EOL> def tearDown ( self ) : <EOL> for fpath in glob ( os . path . join ( DEFAULT_HOME , '<STR_LIT>' ) ) : <EOL> try : <EOL> os . remove ( fpath ) <EOL> except OSError : <EOL> pass <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
