Rough separation of modules and added top level tests to main logic method

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)

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)

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()

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")