Repository: DataBiosphere/toil

Issue:
Write unittest for hot deployment
Ideally should test with and without encryption.



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,9 @@
     author='Benedict Paten',
     author_email='benedict@soe.usc.edu',
     url="https://github.com/BD2KGenomics/toil",
-    install_requires=['bd2k-python-lib>=1.7.dev1'],
+    install_requires=[ 'bd2k-python-lib>=1.7.dev1' ],
+    tests_require=[ 'mock==1.0.1' ],
+    test_suite = 'toil',
     extras_require={
         'mesos': [
             'mesos.interface==0.22.0',
@@ -23,4 +25,4 @@
             'toilStats = toil.utils.toilStats:main',
             'toilRestarts = toil.utils.toilRestarts:main',
             'multijob = toil.batchSystems.multijob:main',
-            'toil-mesos-executor = toil.batchSystems.mesos.executor:main [mesos]'] } )
+            'toil-mesos-executor = toil.batchSystems.mesos.executor:main [mesos]' ] } )
diff --git a/src/toil/job.py b/src/toil/job.py
--- a/src/toil/job.py
+++ b/src/toil/job.py
@@ -81,7 +81,9 @@ def __init__(self, memory=sys.maxint, cpu=sys.maxint, disk=sys.maxint):
         #A follow-on, service or child of a job A, is a "successor" of A, if B
         #is a successor of A, then A is a predecessor of B. 
         self._predecessors = set()
-        #Variables used for serialisation
+        # Note that self.__module__ is not necessarily this module, i.e. job.py. It is the module
+        # defining the class self is an instance of, which may be a subclass of Job that may be
+        # defined in a different module.
         self.userModule = ModuleDescriptor.forModule(self.__module__)
         #See Job.rv()
         self._rvs = {}
@@ -583,7 +585,8 @@ def _makeJobWrappers(self, jobStore, jobsToUUIDs, jobsToJobs, predecessor, rootJ
             with jobStore.writeFileStream(rootJob.jobStoreID) as (fileHandle, fileStoreID):
                 cPickle.dump(self, fileHandle, cPickle.HIGHEST_PROTOCOL)
             jobClassName = self.__class__.__name__
-            jobWrapper.command = ' '.join( ('scriptTree', fileStoreID, jobClassName) + self.userModule)
+            jobWrapper.command = ' '.join( ('scriptTree', fileStoreID, jobClassName)
+                                           + self.userModule.globalize())
             #Update the status of the job on disk
             jobStore.update(jobWrapper)
         else:
@@ -636,7 +639,7 @@ def _serialiseFirstJob(self, jobStore):
             cPickle.dump(self, f, cPickle.HIGHEST_PROTOCOL)
         #Make the first job
         jobClassName = self.__class__.__name__
-        command = ('scriptTree', sharedJobFile, jobClassName) + self.userModule
+        command = ('scriptTree', sharedJobFile, jobClassName) + self.userModule.globalize()
         jobWrapper = self._createEmptyJobForJob(jobStore, command=' '.join( command ))
         #Store the name of the first job in a file in case of restart
         with jobStore.writeSharedFileStream("rootJobStoreID") as f:
@@ -928,7 +931,7 @@ def __init__(self, userFunction, *args, **kwargs):
         disk = kwargs.pop("disk") if "disk" in kwargs else sys.maxint
         memory = kwargs.pop("memory") if "memory" in kwargs else sys.maxint
         Job.__init__(self, memory=memory, cpu=cpu, disk=disk)
-        self.userFunctionModule = ModuleDescriptor.forModule(userFunction.__module__)
+        self.userFunctionModule = ModuleDescriptor.forModule(userFunction.__module__).globalize()
         self.userFunctionName = str(userFunction.__name__)
         self._args=args
         self._kwargs=kwargs
diff --git a/src/toil/resource.py b/src/toil/resource.py
--- a/src/toil/resource.py
+++ b/src/toil/resource.py
@@ -61,7 +61,7 @@ def create( cls, jobStore, leaderPath ):
         contentHash = hashlib.md5( )
         # noinspection PyProtectedMember
         with subcls._load( leaderPath ) as src:
-            with jobStore.writeSharedFileStream( pathHash, isProtected=False) as dst:
+            with jobStore.writeSharedFileStream( pathHash, isProtected=False ) as dst:
                 userScript = src.read( )
                 contentHash.update( userScript )
                 dst.write( userScript )
@@ -113,7 +113,7 @@ def lookup( cls, leaderPath ):
             assert self.pathHash == pathHash
             return self
 
-    def download( self ):
+    def download( self, callback=None ):
         """
         Downloads this resource from its URL to a file on the local system. This method should
         only be invoked on a worker node after the node was setup for accessing resources via
@@ -123,6 +123,8 @@ def download( self ):
         if not os.path.exists( dirPath ):
             tempDirPath = mkdtemp( dir=os.path.dirname( dirPath ), prefix=self.contentHash + "-" )
             self._save( tempDirPath )
+            if callback is not None:
+                callback( tempDirPath )
             try:
                 os.rename( tempDirPath, dirPath )
             except OSError as e:
@@ -144,8 +146,7 @@ def localPath( self ):
     @property
     def localDirPath( self ):
         """
-        The path to the directory containing the resource on the worker. For directory resources
-        this is the same as the localPath property.
+        The path to the directory containing the resource on the worker.
         """
         rootDirPath = os.environ[ self.rootDirPathEnvName ]
         return os.path.join( rootDirPath, self.contentHash )
@@ -278,7 +279,7 @@ def forModule( cls, moduleName ):
         """
         Return an instance of this class representing the module of the given name. If the given
         module name is "__main__", it will be translated to the actual file name of the top-level
-        script without the .py or .pyc extension. This method expects that the module with the
+        script without the .py or .pyc extension. This method assumes that the module with the
         specified name has already been loaded.
         """
         module = sys.modules[ moduleName ]
@@ -294,19 +295,12 @@ def forModule( cls, moduleName ):
                 assert dirPathTail == package
         return cls( dirPath=os.path.sep.join( dirPath ), name=moduleName, extension=extension )
 
-    @classmethod
-    def forDirPath( cls, userModuleDirPath, userModuleName ):
-        pass
-
     @property
     def belongsToToil( self ):
         """
         True if this module is part of the Toil distribution
         """
-        # FIXME: Forcing this to False for now to give this code more exposure but once we
-        # FIXME: hot-deploy toil itself it should be used to shortcut hot-deployment for user
-        # FIXME: scripts inside the toil distribution.
-        return False and self.name.startswith( 'toil.' )
+        return self.name.startswith( 'toil.' )
 
     @property
     def filePath( self ):
@@ -338,10 +332,26 @@ def localize( self ):
             log.warn( "Can't localize module %r", self )
             return self
         else:
-            resource.download( )
+            def stash( tmpDirPath ):
+                # Save the original dirPath such that we can restore it in globalize()
+                with open( os.path.join( tmpDirPath, '.original' ), 'w' ) as f:
+                    f.write( json.dumps( self ) )
+
+            resource.download( callback=stash )
             return self.__class__( dirPath=resource.localDirPath, name=self.name,
                                    extension=self.extension )
 
+    def globalize( self ):
+        try:
+            with open( os.path.join( self.dirPath, '.original' ) ) as f:
+                return self.__class__( *json.loads( f.read( ) ) )
+        except IOError as e:
+            if e.errno == errno.ENOENT:
+                log.warn( "Can't globalize module %r.", self )
+                return self
+            else:
+                raise
+
     @property
     def _resourcePath( self ):
         """


Test patch:
diff --git a/src/toil/test/src/resourceTest.py b/src/toil/test/src/resourceTest.py
new file mode 100644
--- /dev/null
+++ b/src/toil/test/src/resourceTest.py
@@ -0,0 +1,83 @@
+from __future__ import absolute_import
+
+import sys
+
+from mock import MagicMock, patch
+
+from toil.resource import ModuleDescriptor, Resource
+from toil.test import ToilTest
+
+
+class ResourceTest( ToilTest ):
+    def test( self ):
+        # Create a ModuleDescriptor for the module containing ModuleDescriptor, i.e. toil.resource
+        module_name = ModuleDescriptor.__module__
+        self.assertEquals( module_name, 'toil.resource' )
+        module = ModuleDescriptor.forModule( module_name )
+
+        # Assert basic attributes and properties
+        self.assertTrue( module.belongsToToil )
+        self.assertEquals( module.name, module_name )
+        self.assertEquals( module.filePath, sys.modules[ module_name ].__file__ )
+        self.assertTrue( module.filePath.startswith( module.dirPath ) )
+        self.assertTrue( module.dirPath.endswith( '/src' ) )
+
+        # Before the module is saved as a resource, localize() and globalize() are identity
+        # methods. This should log warnings.
+        self.assertIs( module.localize( ), module )
+        self.assertIs( module.globalize( ), module )
+
+        # Create a mock job store ...
+        jobStore = MagicMock( )
+        # ... to generate a fake URL for the resource ...
+        url = 'file://foo.zip'
+        jobStore.getSharedPublicUrl.return_value = url
+        # ... and save the resource to it.
+        resource = module.saveAsResourceTo( jobStore )
+        # Ensure that the URL generation method is actually called, ...
+        jobStore.getSharedPublicUrl.assert_called_once_with( resource.pathHash )
+        # ... and that ensure that writeSharedFileStream is called.
+        jobStore.writeSharedFileStream.assert_called_once_with( resource.pathHash,
+                                                                isProtected=False )
+
+        # Now it gets a bit complicated: Ensure that the context manager returned by the
+        # jobStore's writeSharedFileStream() method is entered and that the file handle yielded
+        # by the context manager is written to once with the zipped source tree from which
+        # 'toil.resource' was orginally imported. Keep the zipped tree around such that we can
+        # mock the download later.
+        file_handle = jobStore.writeSharedFileStream.return_value.__enter__.return_value
+        # The first 0 index selects the first call of write(), the second 0 selects positional
+        # instead of keyword arguments, and the third 0 selects the first positional, i.e. the
+        # contents. This is a bit brittle since it assumes that all the data is written in a
+        # single call to write(). If more calls are made we can easily concatenate them.
+        zipFile = file_handle.write.call_args_list[ 0 ][ 0 ][ 0 ]
+        self.assertTrue( zipFile.startswith( 'PK' ) ) # the magic header for ZIP files
+        self.assertEquals( resource.url, url )
+
+        # Now we're on the worker. Prepare the storage for localized resources
+        Resource.prepareSystem( )
+        # Register the resource for subsequent lookup.
+        resource.register( )
+        # Lookup the resource and ensure that the result is equal to but not the same as the
+        # original resource. Lookup will also be used when we localize the module that was
+        # originally used to create the resource.
+        localResource = Resource.lookup( module._resourcePath )
+        self.assertEquals( resource, localResource )
+        self.assertIsNot( resource, localResource )
+
+        # Now show that we can localize the module using the registered resource. Set up a mock
+        # urlopen() that yields the zipped tree ...
+        mock_urlopen = MagicMock( )
+        mock_urlopen.return_value.read.return_value = zipFile
+        with patch( 'toil.resource.urlopen', mock_urlopen ):
+            # ... and use it to download and unpack the resource
+            localModule = module.localize( )
+        # Name and extension should be equal between original and localized resource ...
+        self.assertEquals( module.name, localModule.name )
+        self.assertEquals( module.extension, localModule.extension )
+        # ... but the directory should be different.
+        self.assertNotEquals( module.dirPath, localModule.dirPath )
+
+        # Show that we can 'undo' localization. This is necessary when the user script's jobs are
+        # invoked on the worker where they generate more child jobs.
+        self.assertEquals( localModule.globalize( ), module )