From 5a38b8f9d5bc80a74ac71ed4b73a1fcf0ee4c82d Mon Sep 17 00:00:00 2001
From: Jim Madge <jmadge@turing.ac.uk>
Date: Thu, 18 Apr 2024 11:57:03 +0100
Subject: [PATCH] Add setitem method

---
 data_safe_haven/config/pulumi.py | 11 +++++++++++
 tests/config/test_pulumi.py      | 16 ++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/data_safe_haven/config/pulumi.py b/data_safe_haven/config/pulumi.py
index 047564af54..2c3bb1bbab 100644
--- a/data_safe_haven/config/pulumi.py
+++ b/data_safe_haven/config/pulumi.py
@@ -45,6 +45,17 @@ def __getitem__(self, key: str):
         msg = f"No configuration for Pulumi stack {key}."
         raise IndexError(msg)
 
+    def __setitem__(self, key: str, value: PulumiStack):
+        if not isinstance(key, str):
+            msg = "'key' must be a string."
+            raise TypeError(msg)
+
+        if key in self.stack_names:
+            msg = f"Stack {key} already exists."
+            raise ValueError(msg)
+
+        self.stacks.append(value)
+
     def __delitem__(self, key: str):
         if not isinstance(key, str):
             msg = "'key' must be a string."
diff --git a/tests/config/test_pulumi.py b/tests/config/test_pulumi.py
index a31ebcba6c..e2e5a81c1c 100644
--- a/tests/config/test_pulumi.py
+++ b/tests/config/test_pulumi.py
@@ -87,6 +87,22 @@ def test_delitem_index_error(self, pulumi_config):
         with raises(IndexError, match="No configuration for Pulumi stack Ringo."):
             del pulumi_config["Ringo"]
 
+    def test_setitem(self, pulumi_config, pulumi_stack):
+        del pulumi_config["my_stack"]
+        assert len(pulumi_config.stack_names) == 1
+        assert "my_stack" not in pulumi_config.stack_names
+        pulumi_config["my_stack"] = pulumi_stack
+        assert len(pulumi_config.stack_names) == 2
+        assert "my_stack" in pulumi_config.stack_names
+
+    def test_setitem_type_error(self, pulumi_config):
+        with raises(TypeError, match="'key' must be a string."):
+            pulumi_config[1] = 5
+
+    def test_setitem_value_error(self, pulumi_config):
+        with raises(ValueError, match="Stack other_stack already exists."):
+            pulumi_config["other_stack"] = 5
+
     def test_stack_names(self, pulumi_config):
         assert "my_stack" in pulumi_config.stack_names