import contextlib

import functools

import os.path

import shlex

import shutil

import subprocess

from ellar.utils.importer import get_main_directory_by_stack

SAMPLE_DUMB_DIRS = get_main_directory_by_stack("__main__/dumbs/", stack_level=1)

SAMPLE_DIRS = get_main_directory_by_stack(

"__main__/test_migrations/samples", stack_level=1

)

def run_command(cmd, **kwargs):

os.chdir(SAMPLE_DIRS)

kwargs.setdefault("stdout", subprocess.PIPE)

kwargs.setdefault("stderr", subprocess.PIPE)

args = shlex.split(f"python {cmd}")

return subprocess.run(args, **kwargs)

def clean_directory(directory):

def _decorator(f):

@functools.wraps(f)

def _wrapper(*args, **kwargs):

try:

f(*args, **kwargs)

finally:

try:

shutil.rmtree(os.path.join(SAMPLE_DUMB_DIRS, directory))

except OSError:

pass

return _wrapper

return _decorator

@contextlib.contextmanager

def set_env_variable(key, value):

old_environ = dict(os.environ)

os.environ[key] = value

yield

os.environ.clear()

os.environ.update(old_environ)