httpie-cli/tests/prompt/base.py
Mickaël Schoentgen 279e387d86 WIP
2021-10-08 10:45:49 +02:00

60 lines
1.8 KiB
Python

import os
import shutil
import sys
import tempfile
import unittest
class TempAppDirTestCase(unittest.TestCase):
"""Set up temporary app data and config directories before every test
method, and delete them afterwards.
"""
def setUp(self):
# Create a temp dir that will contain data and config directories
self.temp_dir = tempfile.mkdtemp()
if sys.platform == 'win32':
self.homes = {
# subdir_name: envvar_name
'data': 'LOCALAPPDATA',
'config': 'LOCALAPPDATA'
}
else:
self.homes = {
# subdir_name: envvar_name
'data': 'XDG_DATA_HOME',
'config': 'XDG_CONFIG_HOME'
}
# Used to restore
self.orig_envvars = {}
for subdir_name, envvar_name in self.homes.items():
if envvar_name in os.environ:
self.orig_envvars[envvar_name] = os.environ[envvar_name]
os.environ[envvar_name] = os.path.join(self.temp_dir, subdir_name)
def tearDown(self):
# Restore envvar values
for name in self.homes.values():
if name in self.orig_envvars:
os.environ[name] = self.orig_envvars[name]
else:
del os.environ[name]
shutil.rmtree(self.temp_dir)
def make_tempfile(self, data='', subdir_name=''):
"""Create a file under self.temp_dir and return the path."""
full_tempdir = os.path.join(self.temp_dir, subdir_name)
if not os.path.exists(full_tempdir):
os.makedirs(full_tempdir)
if isinstance(data, str):
data = data.encode()
with tempfile.NamedTemporaryFile(dir=full_tempdir, delete=False) as f:
f.write(data)
return f.name