[autotest] Factor out more Container code.

Move more code from ContainerBucket into Container.
Add tests.

BUG=chromium:720219

TEST=sudo python site_utils/lxc/container_unittest.py
TEST=sudo python site_utils/lxc/container_bucket_unittest.py
TEST=sudo python site_utils/lxc/lxc_functional_test.py

Change-Id: I76fc3ef864b98c30ce19257d3bd99956bce663dc
Reviewed-on: https://chromium-review.googlesource.com/571882
Commit-Ready: Ben Kwa <kenobi@chromium.org>
Tested-by: Ben Kwa <kenobi@chromium.org>
Reviewed-by: Dan Shi <dshi@google.com>
diff --git a/site_utils/lxc/config.py b/site_utils/lxc/config.py
index 34ae2f9..e071273 100644
--- a/site_utils/lxc/config.py
+++ b/site_utils/lxc/config.py
@@ -95,6 +95,7 @@
 from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import global_config
 from autotest_lib.client.common_lib import utils
+from autotest_lib.site_utils.lxc import constants
 from autotest_lib.site_utils.lxc import utils as lxc_utils
 
 
@@ -108,8 +109,6 @@
 # A temp folder used to store files to be appended to the files inside
 # container.
 APPEND_FOLDER = 'usr/local/ssp_append'
-# Path to folder that contains autotest code inside container.
-CONTAINER_AUTOTEST_DIR = '/usr/local/autotest'
 
 DeployConfig = collections.namedtuple(
         'DeployConfig', ['source', 'target', 'append', 'permission'])
@@ -303,7 +302,7 @@
            made in the host.
 
         """
-        shadow_config = os.path.join(CONTAINER_AUTOTEST_DIR,
+        shadow_config = os.path.join(constants.CONTAINER_AUTOTEST_DIR,
                                      'shadow_config.ini')
 
         # Inject "AUTOSERV/enable_master_ssh: False" in shadow config as
diff --git a/site_utils/lxc/constants.py b/site_utils/lxc/constants.py
index 9fbe485..95a7626 100644
--- a/site_utils/lxc/constants.py
+++ b/site_utils/lxc/constants.py
@@ -7,18 +7,21 @@
 import common
 from autotest_lib.client.bin import utils as common_utils
 from autotest_lib.client.common_lib.global_config import global_config
-from autotest_lib.site_utils.lxc import config as lxc_config
 
 
 # Name of the base container.
 BASE = global_config.get_config_value('AUTOSERV', 'container_base_name')
+
+# Path to folder that contains autotest code inside container.
+CONTAINER_AUTOTEST_DIR = '/usr/local/autotest'
+
 # Naming convention of test container, e.g., test_300_1422862512_2424, where:
 # 300:        The test job ID.
 # 1422862512: The tick when container is created.
 # 2424:       The PID of autoserv that starts the container.
 TEST_CONTAINER_NAME_FMT = 'test_%s_%d_%d'
 # Naming convention of the result directory in test container.
-RESULT_DIR_FMT = os.path.join(lxc_config.CONTAINER_AUTOTEST_DIR, 'results',
+RESULT_DIR_FMT = os.path.join(CONTAINER_AUTOTEST_DIR, 'results',
                               '%s')
 # Attributes to retrieve about containers.
 ATTRIBUTES = ['name', 'state']
@@ -45,7 +48,7 @@
 
 # Path to drone_temp folder in the container, which stores the control file for
 # test job to run.
-CONTROL_TEMP_PATH = os.path.join(lxc_config.CONTAINER_AUTOTEST_DIR, 'drone_tmp')
+CONTROL_TEMP_PATH = os.path.join(CONTAINER_AUTOTEST_DIR, 'drone_tmp')
 
 # Bash command to return the file count in a directory. Test the existence first
 # so the command can return an error code if the directory doesn't exist.
@@ -55,14 +58,18 @@
 APPEND_CMD_FMT = ('echo \'%(content)s\' | sudo tee --append %(file)s'
                   '> /dev/null')
 
-# Path to site-packates in Moblab
-MOBLAB_SITE_PACKAGES = '/usr/lib64/python2.7/site-packages'
-MOBLAB_SITE_PACKAGES_CONTAINER = '/usr/local/lib/python2.7/dist-packages/'
-
 # Flag to indicate it's running in a Moblab. Due to crbug.com/457496, lxc-ls has
 # different behavior in Moblab.
 IS_MOBLAB = common_utils.is_moblab()
 
+if IS_MOBLAB:
+    SITE_PACKAGES_PATH = '/usr/lib64/python2.7/site-packages'
+    CONTAINER_SITE_PACKAGES_PATH = '/usr/local/lib/python2.7/dist-packages/'
+else:
+    SITE_PACKAGES_PATH = os.path.join(common.autotest_dir, 'site-packages')
+    CONTAINER_SITE_PACKAGES_PATH = os.path.join(CONTAINER_AUTOTEST_DIR,
+                                                'site-packages')
+
 # TODO(dshi): If we are adding more logic in how lxc should interact with
 # different systems, we should consider code refactoring to use a setting-style
 # object to store following flags mapping to different systems.
diff --git a/site_utils/lxc/container.py b/site_utils/lxc/container.py
index 4c3e2a4..eedc5e5 100644
--- a/site_utils/lxc/container.py
+++ b/site_utils/lxc/container.py
@@ -10,7 +10,6 @@
 import common
 from autotest_lib.client.bin import utils
 from autotest_lib.client.common_lib import error
-from autotest_lib.site_utils.lxc import config as lxc_config
 from autotest_lib.site_utils.lxc import constants
 from autotest_lib.site_utils.lxc import lxc
 from autotest_lib.site_utils.lxc import utils as lxc_utils
@@ -93,7 +92,7 @@
 
 
     @classmethod
-    def clone(cls, src, new_name, new_path=None, snapshot=False):
+    def clone(cls, src, new_name, new_path=None, snapshot=False, cleanup=False):
         """Creates a clone of this container.
 
         @param src: The original container.
@@ -108,6 +107,23 @@
         if new_path is None:
             new_path = src.container_path
 
+        # If a container exists at this location, clean it up first
+        container_folder = os.path.join(new_path, new_name)
+        if lxc_utils.path_exists(container_folder):
+            if not cleanup:
+                raise error.ContainerError('Container %s already exists.' %
+                                           new_name)
+            container = Container.createFromExistingDir(new_path, new_name)
+            try:
+                container.destroy()
+            except error.CmdError as e:
+                # The container could be created in a incompleted state. Delete
+                # the container folder instead.
+                logging.warn('Failed to destroy container %s, error: %s',
+                             new_name, e)
+                utils.run('sudo rm -rf "%s"' % container_folder)
+
+        # Create and return the new container.
         return cls(new_path, new_name, {}, src, snapshot)
 
 
@@ -293,15 +309,10 @@
         """
         # Test autotest code is setup by verifying a list of
         # (directory, minimum file count)
-        if constants.IS_MOBLAB:
-            site_packages_path = constants.MOBLAB_SITE_PACKAGES_CONTAINER
-        else:
-            site_packages_path = os.path.join(lxc_config.CONTAINER_AUTOTEST_DIR,
-                                              'site-packages')
         directories_to_check = [
-                (lxc_config.CONTAINER_AUTOTEST_DIR, 3),
+                (constants.CONTAINER_AUTOTEST_DIR, 3),
                 (constants.RESULT_DIR_FMT % job_folder, 0),
-                (site_packages_path, 3)]
+                (constants.CONTAINER_SITE_PACKAGES_PATH, 3)]
         for directory, count in directories_to_check:
             result = self.attach_run(command=(constants.COUNT_FILE_CMD %
                                               {'dir': directory})).stdout
@@ -352,3 +363,42 @@
         """Returns whether or not this container is currently running."""
         self.refresh_status()
         return self.state == 'RUNNING'
+
+
+    def set_hostname(self, hostname):
+        """Sets the hostname within the container.  This needs to be called
+        prior to starting the container.
+        """
+        config_file = os.path.join(self.container_path, self.name, 'config')
+        lxc_utsname_setting = (
+                'lxc.utsname = ' +
+                constants.CONTAINER_UTSNAME_FORMAT % hostname)
+        utils.run(
+            constants.APPEND_CMD_FMT % {'content': lxc_utsname_setting,
+                                        'file': config_file})
+
+
+    def install_ssp(self, ssp_url):
+        """Downloads and installs the given server package.
+
+        @param ssp_url: The URL of the ssp to download and install.
+        """
+        usr_local_path = os.path.join(self.rootfs, 'usr', 'local')
+        autotest_pkg_path = os.path.join(usr_local_path,
+                                         'autotest_server_package.tar.bz2')
+        # sudo is required so os.makedirs may not work.
+        utils.run('sudo mkdir -p %s'% usr_local_path)
+
+        lxc.download_extract(ssp_url, autotest_pkg_path, usr_local_path)
+
+
+    def install_control_file(self, control_file):
+        """Installs the given control file.
+        The given file will be moved into the container.
+
+        @param control_file: Path to the control file to install.
+        """
+        dst_path = os.path.join(self.rootfs,
+                                constants.CONTROL_TEMP_PATH.lstrip(os.path.sep))
+        utils.run('sudo mkdir -p %s' % dst_path)
+        utils.run('sudo mv %s %s' % (control_file, dst_path))
diff --git a/site_utils/lxc/container_bucket.py b/site_utils/lxc/container_bucket.py
index 2b6bc64..1ea7e86 100644
--- a/site_utils/lxc/container_bucket.py
+++ b/site_utils/lxc/container_bucket.py
@@ -120,23 +120,23 @@
                         disable_snapshot_clone)
 
         try:
-            return self.clone_container(path=self.container_path,
-                                        name=constants.BASE,
-                                        new_path=self.container_path,
-                                        new_name=name,
-                                        snapshot=use_snapshot,
-                                        cleanup=force_cleanup)
+            return Container.clone(src=self.base_container,
+                                   new_name=name,
+                                   new_path=self.container_path,
+                                   snapshot=use_snapshot,
+                                   cleanup=force_cleanup)
         except error.CmdError:
+            logging.debug('Creating snapshot clone failed. Attempting without '
+                           'snapshot...')
             if not use_snapshot:
                 raise
             else:
                 # Snapshot clone failed, retry clone without snapshot.
-                container = self.clone_container(path=self.container_path,
-                                                 name=constants.BASE,
-                                                 new_path=self.container_path,
-                                                 new_name=name,
-                                                 snapshot=False,
-                                                 cleanup=force_cleanup)
+                container = Container.clone(src=self.base_container,
+                                            new_name=name,
+                                            new_path=self.container_path,
+                                            snapshot=False,
+                                            cleanup=force_cleanup)
                 # Report metadata about retry success.
                 autotest_es.post(
                     use_http=True,
@@ -147,44 +147,6 @@
                 return container
 
 
-    def clone_container(self, path, name, new_path, new_name, snapshot=False,
-                        cleanup=False):
-        """Clone one container from another.
-
-        @param path: LXC path for the source container.
-        @param name: Name of the source container.
-        @param new_path: LXC path for the cloned container.
-        @param new_name: Name for the cloned container.
-        @param snapshot: Whether to snapshot, or create a full clone.
-        @param cleanup: If a container with the given name and path already
-                exist, clean it up first.
-
-        @return: A Container object for the created container.
-
-        @raise ContainerError: If the container already exists.
-        @raise error.CmdError: If lxc-clone call failed for any reason.
-        """
-        # Cleanup existing container with the given name.
-        container_folder = os.path.join(new_path, new_name)
-
-        if lxc_utils.path_exists(container_folder):
-            if not cleanup:
-                raise error.ContainerError('Container %s already exists.' %
-                                           new_name)
-            container = Container.createFromExistingDir(new_path, new_name)
-            try:
-                container.destroy()
-            except error.CmdError as e:
-                # The container could be created in a incompleted state. Delete
-                # the container folder instead.
-                logging.warn('Failed to destroy container %s, error: %s',
-                             name, e)
-                utils.run('sudo rm -rf "%s"' % container_folder)
-
-        lxc_utils.clone(path, name, new_path, new_name, snapshot)
-        return self.get(new_name)
-
-
     @cleanup_if_fail()
     def setup_base(self, name=constants.BASE, force_delete=False):
         """Setup base container.
@@ -344,56 +306,30 @@
         # id and timestamp. For better result view, the container's hostname is
         # set to be a string containing the dut hostname.
         if dut_name:
-            config_file = os.path.join(container.container_path, name, 'config')
-            lxc_utsname_setting = (
-                'lxc.utsname = ' +
-                (constants.CONTAINER_UTSNAME_FORMAT %
-                 dut_name.replace('.', '-')))
-            utils.run(
-                constants.APPEND_CMD_FMT % {'content': lxc_utsname_setting,
-                                            'file': config_file})
+            container.set_hostname(dut_name.replace('.', '-'))
 
         # Deploy server side package
-        usr_local_path = os.path.join(container.rootfs, 'usr', 'local')
-        autotest_pkg_path = os.path.join(usr_local_path,
-                                         'autotest_server_package.tar.bz2')
-        autotest_path = os.path.join(usr_local_path, 'autotest')
-        # sudo is required so os.makedirs may not work.
-        utils.run('sudo mkdir -p %s'% usr_local_path)
+        container.install_ssp(server_package_url)
 
-        lxc.download_extract(
-            server_package_url, autotest_pkg_path, usr_local_path)
         deploy_config_manager = lxc_config.DeployConfigManager(container)
         deploy_config_manager.deploy_pre_start()
 
         # Copy over control file to run the test job.
         if control:
-            container_drone_temp = os.path.join(autotest_path, 'drone_tmp')
-            utils.run('sudo mkdir -p %s'% container_drone_temp)
-            container_control_file = os.path.join(
-                    container_drone_temp, control_file_name)
-            # Move the control file stored in the result folder to container.
-            utils.run('sudo mv %s %s' % (safe_control, container_control_file))
+            container.install_control_file(safe_control)
 
-        if constants.IS_MOBLAB:
-            site_packages_path = constants.MOBLAB_SITE_PACKAGES
-            site_packages_container_path = (
-                constants.MOBLAB_SITE_PACKAGES_CONTAINER[1:])
-        else:
-            site_packages_path = os.path.join(common.autotest_dir,
-                                              'site-packages')
-            site_packages_container_path = os.path.join(
-                    lxc_config.CONTAINER_AUTOTEST_DIR, 'site-packages')
-        mount_entries = [(site_packages_path, site_packages_container_path,
+        mount_entries = [(constants.SITE_PACKAGES_PATH,
+                          constants.CONTAINER_SITE_PACKAGES_PATH,
                           True),
                          (os.path.join(common.autotest_dir, 'puppylab'),
-                          os.path.join(lxc_config.CONTAINER_AUTOTEST_DIR,
+                          os.path.join(constants.CONTAINER_AUTOTEST_DIR,
                                        'puppylab'),
                           True),
                          (result_path,
                           os.path.join(constants.RESULT_DIR_FMT % job_folder),
                           False),
                         ]
+
         for mount_config in deploy_config_manager.mount_configs:
             mount_entries.append((mount_config.source, mount_config.target,
                                   mount_config.readonly))
@@ -404,6 +340,9 @@
         # Update file permissions.
         # TODO(dshi): crbug.com/459344 Skip following action when test container
         # can be unprivileged container.
+        autotest_path = os.path.join(
+                container.rootfs,
+                constants.CONTAINER_AUTOTEST_DIR.lstrip(os.path.sep))
         utils.run('sudo chown -R root "%s"' % autotest_path)
         utils.run('sudo chgrp -R root "%s"' % autotest_path)
 
diff --git a/site_utils/lxc/container_unittest.py b/site_utils/lxc/container_unittest.py
index 922fda3..7aebfd8 100644
--- a/site_utils/lxc/container_unittest.py
+++ b/site_utils/lxc/container_unittest.py
@@ -10,6 +10,7 @@
 import shutil
 import sys
 import unittest
+from contextlib import contextmanager
 
 import common
 from autotest_lib.client.common_lib import error
@@ -77,6 +78,99 @@
                 container.refresh_status()
 
 
+    def testDefaultHostname(self):
+        """Verifies that the zygote starts up with a default hostname that is
+        the lxc container name."""
+        test_name = 'testHostname'
+        with self.createContainer(name=test_name) as container:
+            container.start(wait_for_network=True)
+            hostname = container.attach_run('hostname').stdout.strip()
+            self.assertEqual(test_name, hostname)
+
+
+    @unittest.skip('Setting the container hostname using lxc.utsname does not'
+                   'work on goobuntu.')
+    def testSetHostnameNotRunning(self):
+        """Verifies that the hostname can be set on a stopped container."""
+        with self.createContainer() as container:
+            expected_hostname = 'my-new-hostname'
+            container.set_hostname(expected_hostname)
+            container.start(wait_for_network=True)
+            hostname = container.attach_run('hostname').stdout.strip()
+            self.assertEqual(expected_hostname, hostname)
+
+
+    def testClone(self):
+        """Verifies that cloning a container works as expected."""
+        clone = lxc.Container.clone(src=self.base_container,
+                                    new_name="testClone",
+                                    snapshot=True)
+        try:
+            # Throws an exception if the container is not valid.
+            clone.refresh_status()
+        finally:
+            clone.destroy()
+
+
+    def testCloneWithoutCleanup(self):
+        """Verifies that cloning a container to an existing name will fail as
+        expected.
+        """
+        lxc.Container.clone(src=self.base_container,
+                            new_name="testCloneWithoutCleanup",
+                            snapshot=True)
+        with self.assertRaises(error.ContainerError):
+            lxc.Container.clone(src=self.base_container,
+                                new_name="testCloneWithoutCleanup",
+                                snapshot=True)
+
+
+    def testCloneWithCleanup(self):
+        """Verifies that cloning a container with cleanup works properly."""
+        clone0 = lxc.Container.clone(src=self.base_container,
+                                     new_name="testClone",
+                                     snapshot=True)
+        clone0.start(wait_for_network=False)
+        tmpfile = clone0.attach_run('mktemp').stdout
+        # Verify that our tmpfile exists
+        clone0.attach_run('test -f %s' % tmpfile)
+
+        # Clone another container in place of the existing container.
+        clone1 = lxc.Container.clone(src=self.base_container,
+                                     new_name="testClone",
+                                     snapshot=True,
+                                     cleanup=True)
+        with self.assertRaises(error.CmdError):
+            clone1.attach_run('test -f %s' % tmpfile)
+
+
+    def testInstallControlFile(self):
+        """Verifies that installing a control file in the container works."""
+        _unused, tmpfile = tempfile.mkstemp()
+        with self.createContainer() as container:
+            container.install_control_file(tmpfile)
+            container.start(wait_for_network=False)
+            # Verify that the file is found in the container.
+            container.attach_run(
+                'test -f %s' % os.path.join(lxc.CONTROL_TEMP_PATH,
+                                            os.path.basename(tmpfile)))
+
+
+    @contextmanager
+    def createContainer(self, name=None):
+        """Creates a container from the base container, for testing.
+        Use this to ensure that containers get properly cleaned up after each
+        test.
+
+        @param name: An optional name for the new container.
+        """
+        if name is None:
+            name = self.id().split('.')[-1]
+        container = self.bucket.create_from_base(name)
+        yield container
+        container.destroy()
+
+
 def parse_options():
     """Parse command line inputs.
     """
@@ -90,7 +184,7 @@
     # Hack: python unittest also processes args.  Construct an argv to pass to
     # it, that filters out the options it won't recognize.
     if args.verbose:
-        argv.append('-v')
+        argv.insert(0, '-v')
     argv.insert(0, sys.argv[0])
 
     return args, argv
diff --git a/site_utils/lxc/zygote.py b/site_utils/lxc/zygote.py
index fbb9ee7..b939e54 100644
--- a/site_utils/lxc/zygote.py
+++ b/site_utils/lxc/zygote.py
@@ -7,7 +7,6 @@
 import common
 from autotest_lib.client.bin import utils
 from autotest_lib.site_utils.lxc import Container
-from autotest_lib.site_utils.lxc import config as lxc_config
 from autotest_lib.site_utils.lxc import constants
 from autotest_lib.site_utils.lxc import utils as lxc_utils
 
@@ -50,7 +49,7 @@
             # If creating a new zygote, initialize the host dir.
             if not lxc_utils.path_exists(self.host_path):
                 utils.run('sudo mkdir %s' % self.host_path)
-            self.mount_dir(self.host_path, lxc_config.CONTAINER_AUTOTEST_DIR)
+            self.mount_dir(self.host_path, constants.CONTAINER_AUTOTEST_DIR)
 
 
     def destroy(self, force=True):
@@ -70,13 +69,7 @@
                 'content': '127.0.0.1 %s' % (hostname),
                 'file': '/etc/hosts'})
         else:
-            config_file = os.path.join(self.container_path, self.name, 'config')
-            lxc_utsname_setting = (
-                'lxc.utsname = ' +
-                constants.CONTAINER_UTSNAME_FORMAT % hostname)
-            utils.run(
-                constants.APPEND_CMD_FMT % {'content': lxc_utsname_setting,
-                                            'file': config_file})
+            super(Zygote, self).set_hostname(hostname)
 
 
     def _cleanup_host_mount(self):
diff --git a/site_utils/lxc/zygote_unittest.py b/site_utils/lxc/zygote_unittest.py
index d4ffdc0..8a102b3 100644
--- a/site_utils/lxc/zygote_unittest.py
+++ b/site_utils/lxc/zygote_unittest.py
@@ -15,13 +15,13 @@
 import common
 from autotest_lib.client.bin import utils
 from autotest_lib.site_utils import lxc
-from autotest_lib.site_utils.lxc import config as lxc_config
 from autotest_lib.site_utils.lxc import unittest_logging
 from autotest_lib.site_utils.lxc import utils as lxc_utils
 
 
 options = None
 
+@unittest.skipIf(lxc.IS_MOBLAB, 'Zygotes are not supported on moblab.')
 class ZygoteTests(unittest.TestCase):
     """Unit tests for the Zygote class."""
 
@@ -97,28 +97,6 @@
         # missing.
 
 
-    def testHostname(self):
-        """Verifies that the zygote starts up with a default hostname that is
-        the lxc container name."""
-        test_name = 'testHostname'
-        with self.createZygote(name=test_name) as zygote:
-            zygote.start(wait_for_network=True)
-            hostname = zygote.attach_run('hostname -f').stdout.strip()
-            self.assertEqual(test_name, hostname)
-
-
-    @unittest.skip('Setting the container hostname using lxc.utsname does not'
-                   'work on goobuntu.')
-    def testSetHostnameNotRunning(self):
-        """Verifies that the hostname can be set on a stopped container."""
-        with self.createZygote() as zygote:
-            expected_hostname = 'my-new-hostname'
-            zygote.set_hostname(expected_hostname)
-            zygote.start(wait_for_network=True)
-            hostname = zygote.attach_run('hostname -f').stdout.strip()
-            self.assertEqual(expected_hostname, hostname)
-
-
     def testSetHostnameRunning(self):
         """Verifies that the hostname can be set on a running container."""
         with self.createZygote() as zygote:
@@ -140,7 +118,7 @@
 
             self.verifyBindMount(
                 zygote,
-                container_path=lxc_config.CONTAINER_AUTOTEST_DIR,
+                container_path=lxc.CONTAINER_AUTOTEST_DIR,
                 host_path=zygote.host_path)
 
 
@@ -164,11 +142,11 @@
 
             self.verifyBindMount(
                 zygote,
-                container_path=lxc_config.CONTAINER_AUTOTEST_DIR,
+                container_path=lxc.CONTAINER_AUTOTEST_DIR,
                 host_path=zygote.host_path)
 
             # Verify that the old directory contents was preserved.
-            cmd = 'cat %s' % os.path.join(lxc_config.CONTAINER_AUTOTEST_DIR,
+            cmd = 'cat %s' % os.path.join(lxc.CONTAINER_AUTOTEST_DIR,
                                           test_filename)
             test_output = zygote.attach_run(cmd).stdout.strip()
             self.assertEqual(test_string, test_output)
@@ -243,7 +221,6 @@
 if __name__ == '__main__':
     options, unittest_argv = parse_options()
 
-
     log_level=(logging.DEBUG if options.verbose else logging.INFO)
     unittest_logging.setup(log_level)