From ea9bc944c9456610e7aa13fad9f734582ca7457b Mon Sep 17 00:00:00 2001 From: Simon Arnold Date: Tue, 29 Oct 2019 21:29:27 +0100 Subject: [PATCH] Rough separation of modules and added top level tests to main logic method --- atmos.py | 80 +++++--------------------------------------------- credentials.py | 36 +++++++++++++++++++++++ tests.py | 69 +++++++++++++++++++++++++++++++++++++++++++ workspaces.py | 33 +++++++++++++++++++++ 4 files changed, 146 insertions(+), 72 deletions(-) create mode 100644 credentials.py create mode 100644 tests.py create mode 100644 workspaces.py diff --git a/atmos.py b/atmos.py index db8b98a..a974f08 100755 --- a/atmos.py +++ b/atmos.py @@ -1,6 +1,8 @@ #!/usr/bin/env python3 import argparse, subprocess, shlex, sys, os, glob +import workspaces +import credentials def main(argv): parser = argparse.ArgumentParser(description='Control Terraform Workspaces.') @@ -20,9 +22,9 @@ def determine_actions(args, params): if (is_git_directory()) and not (args.m): if (args.e): aws_creds_file = aws_creds_file + "-atmos" - workspace_manager() + workspaces.workspace_manager() - workspace = get_env() + workspace = workspaces.get_env() workspace_vars = workspace if (args.project): workspace = args.project + "-" + workspace @@ -39,86 +41,20 @@ def determine_actions(args, params): cmd = cmd + ' ' + param if (args.e): - generate_creds(args) + credentials.generate(args) if (args.verbose): print("Atmos will run: " + cmd) - print('Terraform {args} using env vars in {env}'.format(args=args.command, env=workspace_vars)) + run_cmd(cmd) + +def run_cmd(cmd): with subprocess.Popen(shlex.split(cmd)) as proc: exit # Start process but kill py program def is_git_directory(): return subprocess.call(['git', 'branch'], stderr=subprocess.STDOUT, stdout = open(os.devnull, 'w')) == 0 -def workspace_manager(): - branch = subprocess.getoutput("git rev-parse --abbrev-ref HEAD") - if branch == "master": - branch = "default" - else: - if branch not in get_valid_envs(): - branch = "default" - - if get_env() != branch: - print("[INFO]: Terraform workspace & git branch have diverged. Changing workspace to git branch...") - subprocess.call(["terraform", "workspace", "new", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w')) - subprocess.call(["terraform", "workspace", "select", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w')) - -def generate_creds(args): - current_workspace = get_env() - workspaces = ['default'] - - if current_workspace != 'default': - workspaces.append(current_workspace) - - project_name = "" - if (args.project): - delimeter = "-" - if (args.e): - delimeter = "_" - project_name = args.project.upper() + delimeter - - contents = "" - for workspace in workspaces: - access_key_name = project_name + workspace.upper() + '_ACCESS_KEY_ID' - secret_key_name = project_name + workspace.upper() + '_SECRET_ACCESS_KEY' - - if (args.verbose): - print(access_key_name) - print(secret_key_name) - - contents = contents + "[{workspace}]\n".format(workspace=project_name + workspace) - try: - contents = contents + "aws_access_key_id=" + os.environ.get(access_key_name) + "\n" - except: - print("[ERROR]: Env Variable " + access_key_name + " not found.") - sys.exit(1) - try: - contents = contents + "aws_secret_access_key=" + os.environ.get(secret_key_name) + "\n" - except: - print("[ERROR]: Env Variable " + secret_key_name + " not found.") - sys.exit(1) - with open(os.path.expanduser('~/.aws/credentials'), 'w+') as f: - f.write(contents) - -def get_valid_envs(): - try: - # Use var files when present, otherwise default to default - return [os.path.splitext(os.path.basename(x))[0] for x in glob.glob("vars/*.tfvars")] - except FileNotFoundError: - return False - -def get_env(): - try: - tf_env = "" - with open('.terraform/environment', 'r') as f: - tf_env = f.readline() - except: - return("default") - if str(tf_env) in get_valid_envs(): - return(tf_env) - else: - return("default") if __name__ == "__main__": main(sys.argv) \ No newline at end of file diff --git a/credentials.py b/credentials.py new file mode 100644 index 0000000..bd8b2ea --- /dev/null +++ b/credentials.py @@ -0,0 +1,36 @@ +def generate(args): + current_workspace = workspaces.get_env() + workspaces_names = ['default'] + + if current_workspace != 'default': + workspaces_names.append(current_workspace) + + project_name = "" + if (args.project): + delimeter = "-" + if (args.e): + delimeter = "_" + project_name = args.project.upper() + delimeter + + contents = "" + for workspace in workspaces_names: + access_key_name = project_name + workspace.upper() + '_ACCESS_KEY_ID' + secret_key_name = project_name + workspace.upper() + '_SECRET_ACCESS_KEY' + + if (args.verbose): + print(access_key_name) + print(secret_key_name) + + contents = contents + "[{workspace}]\n".format(workspace=project_name + workspace) + try: + contents = contents + "aws_access_key_id=" + os.environ.get(access_key_name) + "\n" + except: + print("[ERROR]: Env Variable " + access_key_name + " not found.") + sys.exit(1) + try: + contents = contents + "aws_secret_access_key=" + os.environ.get(secret_key_name) + "\n" + except: + print("[ERROR]: Env Variable " + secret_key_name + " not found.") + sys.exit(1) + with open(os.path.expanduser('~/.aws/credentials'), 'w+') as f: + f.write(contents) diff --git a/tests.py b/tests.py new file mode 100644 index 0000000..9f1373a --- /dev/null +++ b/tests.py @@ -0,0 +1,69 @@ +import unittest +from unittest.mock import MagicMock, patch +import argparse +import atmos + +class DetermineActionsTests(unittest.TestCase): + + def setUp(self): + self.input_args = argparse.Namespace(command="mytestcommand", e=False, m=False, n=False, project="", verbose=False) + atmos.is_git_directory = MagicMock(return_value=False) + atmos.run_cmd = MagicMock() + + def test_whenCalledWithNoAdditionalArgs_shouldRunTheTerraformCommand(self): + """Simplest case where atmos is only called with a command and no further arguments""" + atmos.determine_actions(self.input_args, []) + atmos.run_cmd.assert_called_with("terraform mytestcommand") + + def test_whenCalledWithParams_theyAreAppended(self): + """Should append all params after the command""" + atmos.determine_actions(self.input_args, ["--myparam", "myvalue"]) + atmos.run_cmd.assert_called_with("terraform mytestcommand --myparam myvalue") + + @patch("workspaces.get_env") + def test_whenCalledWithInitCommand_shouldAppendVarsAndCreds(self, mocked_get_env): + """Case where var-file, -var workpace=xyz and -var shared_credentials is appended""" + self.input_args.command = "init" + mocked_get_env.return_value = "mytestenv" + atmos.determine_actions(self.input_args, []) + atmos.run_cmd.assert_called_with('terraform init -var-file=vars/mytestenv.tfvars -var "workspace=mytestenv" -var "shared_credentials_file=$HOME/.aws/credentials"') + + @patch("workspaces.get_env") + def test_whenCalledWithPlanCommand_shouldAppendVarsAndCreds(self, mocked_get_env): + """Case where var-file, -var workpace=xyz and -var shared_credentials is appended""" + self.input_args.command = "plan" + mocked_get_env.return_value = "mytestenv" + atmos.determine_actions(self.input_args, []) + atmos.run_cmd.assert_called_with('terraform plan -var-file=vars/mytestenv.tfvars -var "workspace=mytestenv" -var "shared_credentials_file=$HOME/.aws/credentials"') + + @patch("workspaces.get_env") + def test_whenCalledWithApplyCommand_shouldAppendVarsAndCreds(self, mocked_get_env): + """Case where var-file, -var workpace=xyz and -var shared_credentials is appended""" + self.input_args.command = "apply" + mocked_get_env.return_value = "mytestenv" + atmos.determine_actions(self.input_args, []) + atmos.run_cmd.assert_called_with('terraform apply -var-file=vars/mytestenv.tfvars -var "workspace=mytestenv" -var "shared_credentials_file=$HOME/.aws/credentials"') + + @patch("workspaces.get_env") + def test_whenCalledWithDestroyCommand_shouldAppendVarsAndCreds(self, mocked_get_env): + """Case where var-file, -var workpace=xyz and -var shared_credentials is appended""" + self.input_args.command = "destroy" + mocked_get_env.return_value = "mytestenv" + atmos.determine_actions(self.input_args, []) + atmos.run_cmd.assert_called_with('terraform destroy -var-file=vars/mytestenv.tfvars -var "workspace=mytestenv" -var "shared_credentials_file=$HOME/.aws/credentials"') + + @patch("workspaces.workspace_manager") + def test_whenInAGitRepo_andManualArgIsNotGiven_andEnvironmentArgIsNotGiven_shouldCallTheWorkspaceManager(self, mocked_workspace_manager): + atmos.is_git_directory.return_value = True + atmos.determine_actions(self.input_args, []) + mocked_workspace_manager.assert_called_once() + + @patch("credentials.generate") + def test_whenEnvironmentArgIsGiven_shouldGenerateCredentials(self, mocked_generate): + self.input_args.e = True + atmos.determine_actions(self.input_args, []) + mocked_generate.assert_called_once() + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/workspaces.py b/workspaces.py new file mode 100644 index 0000000..617eda7 --- /dev/null +++ b/workspaces.py @@ -0,0 +1,33 @@ +import subprocess, os, glob + +def workspace_manager(): + branch = subprocess.getoutput("git rev-parse --abbrev-ref HEAD") + if branch == "master": + branch = "default" + else: + if branch not in get_valid_envs(): + branch = "default" + + if get_env() != branch: + print("[INFO]: Terraform workspace & git branch have diverged. Changing workspace to git branch...") + subprocess.call(["terraform", "workspace", "new", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w')) + subprocess.call(["terraform", "workspace", "select", branch], stderr=subprocess.STDOUT, stdout=open(os.devnull, 'w')) + +def get_valid_envs(): + try: + # Use var files when present, otherwise default to default + return [os.path.splitext(os.path.basename(x))[0] for x in glob.glob("vars/*.tfvars")] + except FileNotFoundError: + return False + +def get_env(): + try: + tf_env = "" + with open('.terraform/environment', 'r') as f: + tf_env = f.readline() + except: + return("default") + if str(tf_env) in get_valid_envs(): + return(tf_env) + else: + return("default") \ No newline at end of file