subiquity/subiquitycore/tests/__init__.py

import functools
import os
import shutil
import tempfile
from unittest import IsolatedAsyncioTestCase


class SubiTestCase(IsolatedAsyncioTestCase):
    def tmp_dir(self, dir=None, cleanup=True):
        # return a full path to a temporary directory that will be cleaned up.
        if dir is None:
            tmpd = tempfile.mkdtemp(prefix="subiquity-%s." % self.__class__.__name__)
        else:
            tmpd = tempfile.mkdtemp(dir=dir)
        self.addCleanup(functools.partial(shutil.rmtree, tmpd))
        return tmpd

    def tmp_path(self, path, dir=None):
        # return an absolute path to 'path' under dir.
        # if dir is None, one will be created with tmp_dir()
        # the file is not created or modified.
        if dir is None:
            dir = self.tmp_dir()
        return os.path.normpath(os.path.abspath(os.path.join(dir, path)))

    def assert_contents(self, path, expected_contents):
        with open(path, "r") as fp:
            self.assertEqual(expected_contents, fp.read())


def populate_dir(path, files):
    if not os.path.exists(path):
        os.makedirs(path)
    ret = []
    for name, content in files.items():
        p = os.path.sep.join([path, name])
        if not os.path.isdir(os.path.dirname(p)):
            os.makedirs(os.path.dirname(p))
        with open(p, "wb") as fp:
            if isinstance(content, bytes):
                fp.write(content)
            else:
                fp.write(content.encode("utf-8"))
            fp.close()
        ret.append(p)
    return ret