Allow passing environment variables to terraform commands

For example, to set the AWS region and profile.

Author: tomduckering

src/infrablocks/invoke_terraform/invoke_executor.py

@@ -1,13 +1,16 @@
-from typing import Iterable
+from typing import Iterable, Optional
 
 from invoke.context import Context
 
 import infrablocks.invoke_terraform.terraform as tf
+from infrablocks.invoke_terraform.terraform.terraform import Environment
 
 
 class InvokeExecutor(tf.Executor):
     def __init__(self, context: Context):
         self._context = context
 
-    def execute(self, command: Iterable[str]) -> None:
-        self._context.run(" ".join(command))
+    def execute(
+        self, command: Iterable[str], env: Optional[Environment] = None
+    ) -> None:
+        self._context.run(" ".join(command), env=(env or {}))

src/infrablocks/invoke_terraform/task_factory.py

@@ -6,6 +6,7 @@
 
 import infrablocks.invoke_factory as invoke_factory
 import infrablocks.invoke_terraform.terraform as tf
+from infrablocks.invoke_terraform.terraform.terraform import Environment
 from infrablocks.invoke_terraform.terraform_factory import TerraformFactory
 
 
@@ -21,6 +22,7 @@ class Configuration:
     variables: tf.Variables
     workspace: Optional[str]
     init_configuration: InitConfiguration
+    environment: Optional[Environment] = None
     auto_approve: bool = True
 
     @staticmethod
@@ -32,6 +34,7 @@ def create_empty():
             init_configuration=InitConfiguration(
                 backend_config={}, reconfigure=False
             ),
+            environment={},
         )
 
 
@@ -74,6 +77,7 @@ def plan(context: Context, arguments: invoke_factory.Arguments):
             terraform.plan(
                 chdir=configuration.source_directory,
                 vars=configuration.variables,
+                environment=configuration.environment,
             )
 
         return plan
@@ -90,6 +94,7 @@ def apply(context: Context, arguments: invoke_factory.Arguments):
                 chdir=configuration.source_directory,
                 vars=configuration.variables,
                 autoapprove=configuration.auto_approve,
+                environment=configuration.environment,
             )
 
         return apply
@@ -111,13 +116,15 @@ def _pre_command_setup(
             chdir=configuration.source_directory,
             backend_config=configuration.init_configuration.backend_config,
             reconfigure=configuration.init_configuration.reconfigure,
+            environment=configuration.environment,
         )
 
         if configuration.workspace is not None:
             terraform.select_workspace(
                 configuration.workspace,
                 chdir=configuration.source_directory,
                 or_create=True,
+                environment=configuration.environment,
             )
 
         return terraform, configuration

src/infrablocks/invoke_terraform/terraform/terraform.py

@@ -3,10 +3,13 @@
 type ConfigurationValue = Union[bool, int, float, str, None]
 type Variables = Dict[str, ConfigurationValue]
 type BackendConfig = Union[str, Dict[str, ConfigurationValue]]
+type Environment = Dict[str, str]
 
 
 class Executor:
-    def execute(self, command: Iterable[str]) -> None:
+    def execute(
+        self, command: Iterable[str], env: Optional[Environment]
+    ) -> None:
         raise Exception("NotImplementedException")
 
 
@@ -19,6 +22,7 @@ def init(
         chdir: Optional[str] = None,
         backend_config: Optional[BackendConfig] = {},
         reconfigure: Optional[bool] = False,
+        environment: Optional[Environment] = None,
     ):
         base_command = self._build_base_command(chdir)
         command = (
@@ -30,21 +34,25 @@ def init(
         if reconfigure:
             command = command + ["-reconfigure"]
 
-        self._executor.execute(command)
+        self._executor.execute(command, env=environment)
 
     def plan(
-        self, chdir: Optional[str] = None, vars: Optional[Variables] = {}
+        self,
+        chdir: Optional[str] = None,
+        vars: Optional[Variables] = {},
+        environment: Optional[Environment] = None,
     ):
         base_command = self._build_base_command(chdir)
         command = base_command + ["plan"] + self._build_vars(vars)
 
-        self._executor.execute(command)
+        self._executor.execute(command, env=environment)
 
     def apply(
         self,
         chdir: Optional[str] = None,
         vars: Optional[Variables] = {},
         autoapprove: bool = False,
+        environment: Optional[Environment] = None,
     ):
         base_command = self._build_base_command(chdir)
         autoapprove_flag = ["-auto-approve"] if autoapprove else []
@@ -55,13 +63,14 @@ def apply(
             + self._build_vars(vars)
         )
 
-        self._executor.execute(command)
+        self._executor.execute(command, env=environment)
 
     def select_workspace(
         self,
         workspace: str,
         chdir: Optional[str] = None,
         or_create: bool = False,
+        environment: Optional[Environment] = None,
     ):
         base_command = self._build_base_command(chdir)
         command = base_command + ["workspace", "select"]
@@ -71,7 +80,7 @@ def select_workspace(
 
         command = command + [workspace]
 
-        self._executor.execute(command)
+        self._executor.execute(command, env=environment)
 
     @staticmethod
     def _build_base_command(chdir: Optional[str]) -> List[str]:

tests/unit/infrablocks/invoke_terraform/terraform/test_terraform.py

@@ -21,7 +21,9 @@ def test_init_executes(self):
 
         terraform.init()
 
-        executor.execute.assert_called_once_with(["terraform", "init"])
+        executor.execute.assert_called_once_with(
+            ["terraform", "init"], env=None
+        )
 
     def test_init_executes_with_chdir(self):
         executor = Mock(spec=Executor)
@@ -30,7 +32,7 @@ def test_init_executes_with_chdir(self):
         terraform.init(chdir="/some/dir")
 
         executor.execute.assert_called_once_with(
-            ["terraform", "-chdir=/some/dir", "init"]
+            ["terraform", "-chdir=/some/dir", "init"], env=None
         )
 
     def test_init_executes_with_backend_config_dictionary(self):
@@ -41,7 +43,7 @@ def test_init_executes_with_backend_config_dictionary(self):
         terraform.init(backend_config=backend_config)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "init", '-backend-config="foo=1"']
+            ["terraform", "init", '-backend-config="foo=1"'], env=None
         )
 
     def test_init_executes_with_backend_config_path(self):
@@ -52,7 +54,8 @@ def test_init_executes_with_backend_config_path(self):
         terraform.init(backend_config=backend_config)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "init", "-backend-config=/some/config.tfvars"]
+            ["terraform", "init", "-backend-config=/some/config.tfvars"],
+            env=None,
         )
 
     def test_init_executes_with_reconfigure(self):
@@ -62,7 +65,18 @@ def test_init_executes_with_reconfigure(self):
         terraform.init(reconfigure=True)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "init", "-reconfigure"]
+            ["terraform", "init", "-reconfigure"], env=None
+        )
+
+    def test_init_executes_with_environment(self):
+        executor = Mock(spec=Executor)
+        terraform = Terraform(executor)
+        environment = {"ENV_VAR": "value"}
+
+        terraform.init(environment=environment)
+
+        executor.execute.assert_called_once_with(
+            ["terraform", "init"], env=environment
         )
 
     def test_plan_executes(self):
@@ -71,7 +85,9 @@ def test_plan_executes(self):
 
         terraform.plan()
 
-        executor.execute.assert_called_once_with(["terraform", "plan"])
+        executor.execute.assert_called_once_with(
+            ["terraform", "plan"], env=None
+        )
 
     def test_plan_executes_with_chdir(self):
         executor = Mock(spec=Executor)
@@ -80,7 +96,7 @@ def test_plan_executes_with_chdir(self):
         terraform.plan(chdir="/some/dir")
 
         executor.execute.assert_called_once_with(
-            ["terraform", "-chdir=/some/dir", "plan"]
+            ["terraform", "-chdir=/some/dir", "plan"], env=None
         )
 
     def test_plan_executes_with_vars(self):
@@ -91,7 +107,18 @@ def test_plan_executes_with_vars(self):
         terraform.plan(vars=variables)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "plan", '-var="foo=1"']
+            ["terraform", "plan", '-var="foo=1"'], env=None
+        )
+
+    def test_plan_executes_with_environment(self):
+        executor = Mock(spec=Executor)
+        terraform = Terraform(executor)
+        environment = {"ENV_VAR": "value"}
+
+        terraform.plan(environment=environment)
+
+        executor.execute.assert_called_once_with(
+            ["terraform", "plan"], env=environment
         )
 
     def test_apply_executes(self):
@@ -100,7 +127,9 @@ def test_apply_executes(self):
 
         terraform.apply()
 
-        executor.execute.assert_called_once_with(["terraform", "apply"])
+        executor.execute.assert_called_once_with(
+            ["terraform", "apply"], env=None
+        )
 
     def test_apply_executes_with_chdir(self):
         executor = Mock(spec=Executor)
@@ -109,7 +138,7 @@ def test_apply_executes_with_chdir(self):
         terraform.apply(chdir="/some/dir")
 
         executor.execute.assert_called_once_with(
-            ["terraform", "-chdir=/some/dir", "apply"]
+            ["terraform", "-chdir=/some/dir", "apply"], env=None
         )
 
     def test_apply_executes_with_vars(self):
@@ -120,7 +149,7 @@ def test_apply_executes_with_vars(self):
         terraform.apply(vars=variables)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "apply", '-var="foo=1"']
+            ["terraform", "apply", '-var="foo=1"'], env=None
         )
 
     def test_apply_executes_with_autoapprove(self):
@@ -130,7 +159,18 @@ def test_apply_executes_with_autoapprove(self):
         terraform.apply(autoapprove=True)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "apply", "-auto-approve"]
+            ["terraform", "apply", "-auto-approve"], env=None
+        )
+
+    def test_apply_executes_with_environment(self):
+        executor = Mock(spec=Executor)
+        terraform = Terraform(executor)
+        environment = {"ENV_VAR": "value"}
+
+        terraform.apply(environment=environment)
+
+        executor.execute.assert_called_once_with(
+            ["terraform", "apply"], env=environment
         )
 
     def test_select_workspace_executes(self):
@@ -141,7 +181,7 @@ def test_select_workspace_executes(self):
         terraform.select_workspace(workspace)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "workspace", "select", workspace]
+            ["terraform", "workspace", "select", workspace], env=None
         )
 
     def test_select_workspace_executes_with_chdir(self):
@@ -152,7 +192,14 @@ def test_select_workspace_executes_with_chdir(self):
         terraform.select_workspace(workspace, chdir="/some/dir")
 
         executor.execute.assert_called_once_with(
-            ["terraform", "-chdir=/some/dir", "workspace", "select", workspace]
+            [
+                "terraform",
+                "-chdir=/some/dir",
+                "workspace",
+                "select",
+                workspace,
+            ],
+            env=None,
         )
 
     def test_select_workspace_executes_with_or_create(self):
@@ -163,5 +210,18 @@ def test_select_workspace_executes_with_or_create(self):
         terraform.select_workspace(workspace, or_create=True)
 
         executor.execute.assert_called_once_with(
-            ["terraform", "workspace", "select", "-or-create=true", workspace]
+            ["terraform", "workspace", "select", "-or-create=true", workspace],
+            env=None,
+        )
+
+    def test_select_workspace_executes_with_environment(self):
+        executor = Mock(spec=Executor)
+        terraform = Terraform(executor)
+        workspace = "workspace"
+        environment = {"ENV_VAR": "value"}
+
+        terraform.select_workspace(workspace, environment=environment)
+
+        executor.execute.assert_called_once_with(
+            ["terraform", "workspace", "select", workspace], env=environment
         )

tests/unit/infrablocks/invoke_terraform/test_invoke_executor.py

@@ -0,0 +1,45 @@
+from unittest.mock import Mock
+
+from invoke.context import Context
+
+from infrablocks.invoke_terraform.invoke_executor import InvokeExecutor
+
+
+class TestInvokeExecutor:
+    def test_run_invoked_with_command(self):
+        context = Mock(spec=Context)
+
+        executor = InvokeExecutor(context)
+
+        executor.execute(["some_command", "arg1", "arg2"])
+
+        context.run.assert_called_once_with(
+            "some_command arg1 arg2",
+            env={},
+        )
+
+    def test_run_invoked_with_env(self):
+        context = Mock(spec=Context)
+
+        executor = InvokeExecutor(context)
+
+        executor.execute(
+            ["some_command", "arg1", "arg2"], env={"ENV_VAR": "value"}
+        )
+
+        context.run.assert_called_once_with(
+            "some_command arg1 arg2",
+            env={"ENV_VAR": "value"},
+        )
+
+    def test_run_invoked_with_empty_env(self):
+        context = Mock(spec=Context)
+
+        executor = InvokeExecutor(context)
+
+        executor.execute(["some_command", "arg1", "arg2"], env=None)
+
+        context.run.assert_called_once_with(
+            "some_command arg1 arg2",
+            env={},
+        )