import logging
from contextlib import contextmanager
from subprocess import CalledProcessError

from bd2k.util.iterables import concat

from toil.test import needs_mesos, ApplianceTestSupport

log = logging.getLogger(__name__)


@needs_mesos
class HotDeploymentTest(ApplianceTestSupport):
    def setUp(self):
        logging.basicConfig(level=logging.INFO)
        super(HotDeploymentTest, self).setUp()

    @contextmanager
    def _venvApplianceCluster(self):
        """
        Creates an appliance cluster with a virtualenv at './venv' on the leader and a temporary
        directory on the host mounted at /data in the leader and worker containers.
        """
        dataDirPath = self._createTempDir(purpose='data')
        with self._applianceCluster(mounts={dataDirPath: '/data'}) as (leader, worker):
            try:
                leader.runOnAppliance('virtualenv',
                                      '--system-site-packages',
                                      '--never-download',
                                      'venv')
                leader.runOnAppliance('venv/bin/pip', 'list') # For diagnostic purposes
                yield leader, worker
            finally:
                # Without this step, we would leak files owned by root on the host's file system
                leader.runOnAppliance('rm', '-rf', '/data/*')

    sitePackages = 'venv/lib/python2.7/site-packages'

    def testRestart(self):
        """
        Test whether hot-deployment works on restart.
        """
        with self._venvApplianceCluster() as (leader, worker):
            def userScript():
                from toil.job import Job
                from toil.common import Toil

                # noinspection PyUnusedLocal
                def job(job, disk='10M', cores=1, memory='10M'):
                    # make job fail during the first invocation
                    assert job.fileStore.jobStore.config.restart

                if __name__ == '__main__':
                    options = Job.Runner.getDefaultArgumentParser().parse_args()
                    with Toil(options) as toil:
                        if toil.config.restart:
                            toil.restart()
                        else:
                            toil.start(Job.wrapJobFn(job))

            leader.deployScript(path=self.sitePackages,
                                packagePath='foo.bar',
                                script=userScript)

            pythonArgs = ['venv/bin/python', '-m', 'foo.bar']
            toilArgs = ['--logDebug',
                        '--batchSystem=mesos',
                        '--mesosMaster=localhost:5050',
                        '--defaultMemory=10M',
                        '/data/jobstore']
            command = concat(pythonArgs, toilArgs)
            self.assertRaises(CalledProcessError, leader.runOnAppliance, *command)
            command = concat(pythonArgs, '--restart', toilArgs)
            leader.runOnAppliance(*command)

@@ -2972,16 +2972,2960 @@
 iance(*command)%0A
+%0A    def testSplitRootPackages(self):%0A        %22%22%22%0A        Test whether hot-deployment works with a virtualenv in which jobs are defined in%0A        completely separate branches of the package hierarchy. Initially, hot deployment did%0A        deploy the entire virtualenv but jobs could only be defined in one branch of the package%0A        hierarchy. We define a branch as the maximum set of fully qualified package paths that%0A        share the same first component. IOW, a.b and a.c are in the same branch, while a.b and%0A        d.c are not.%0A        %22%22%22%0A        with self._venvApplianceCluster() as (leader, worker):%0A%0A            # Deploy the library module with job definitions%0A            def libraryModule():%0A                # noinspection PyUnusedLocal%0A                def libraryJob(job):%0A                    open('/data/foo.txt', 'w').close()%0A%0A            leader.deployScript(path=self.sitePackages,%0A                                packagePath='toil_lib.foo',%0A                                script=libraryModule)%0A%0A            # Deploy the user script%0A            def userScript():%0A                from toil.job import Job%0A                from toil.common import Toil%0A                # noinspection PyUnresolvedReferences%0A                from toil_lib.foo import libraryJob%0A%0A                # noinspection PyUnusedLocal%0A                def job(job, disk='10M', cores=1, memory='10M'):%0A                    # Double the requirements to prevent chaining as chaining might hide problems%0A                    # in hot deployment code.%0A                    job.addChildJobFn(libraryJob, disk='20M', cores=cores, memory=memory)%0A%0A                if __name__ == '__main__':%0A                    options = Job.Runner.getDefaultArgumentParser().parse_args()%0A                    with Toil(options) as toil:%0A                        if toil.config.restart:%0A                            toil.restart()%0A                        else:%0A                            toil.start(Job.wrapJobFn(job))%0A%0A            leader.deployScript(path=self.sitePackages,%0A                                packagePath='toil_script.bar',%0A                                script=userScript)%0A%0A            # Assert that output file isn't there%0A            worker.runOnAppliance('test', '!', '-f', '/data/foo.txt')%0A            # Just being paranoid%0A            self.assertRaises(CalledProcessError,%0A                              worker.runOnAppliance, 'test', '-f', '/data/foo.txt')%0A            leader.runOnAppliance('venv/bin/python',%0A                                  '-m', 'toil_script.bar',%0A                                  '--logDebug',%0A                                  '--batchSystem=mesos',%0A                                  '--mesosMaster=localhost:5050',%0A                                  '--defaultMemory=10M',%0A                                  '/data/jobstore')%0A            # Assert that out output file is there%0A            worker.runOnAppliance('test', '-f', '/data/foo.txt')%0A

Test hot-deploymen of venv with split root packages (connected to #1152)