subiquity/subiquitycore/tests/__init__.py

45 lines
1.3 KiB
Python

import functools
import os
import shutil
import tempfile
from unittest import TestCase
class SubiTestCase(TestCase):
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 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