import unittest

import sys

import tempfile

class TestExitCode(unittest.TestCase):

def setUp(self):

# to be able to import mprof

sys.path.append('.')

from mprof import run_action

self.run_action = run_action

def test_exit_code_success(self):

s = "1+1"

tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')

with tmpfile as ofile:

ofile.write(s)

ofile.flush()

sys.argv = ['<ignored>', '--exit-code', tmpfile.name]

self.assertRaisesRegex(SystemExit, '0', self.run_action)

def test_exit_code_fail(self):

s = "raise RuntimeError('I am not working nicely')"

tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')

with tmpfile as ofile:

ofile.write(s)

ofile.flush()

sys.argv = ['<ignored>', '--exit-code', tmpfile.name]

self.assertRaisesRegex(SystemExit, '1', self.run_action)

def test_no_exit_code_success(self):

s = "raise RuntimeError('I am not working nicely')"

tmpfile = tempfile.NamedTemporaryFile('w', suffix='.py')

with tmpfile as ofile:

ofile.write(s)

ofile.flush()

sys.argv = ['<ignored>', tmpfile.name]

self.run_action()

if __name__ == '__main__':

unittest.main()