make TerraformState classmethods more clear

pytest_terraform/exceptions.py

@@ -12,3 +12,11 @@ class InvalidOption(PytestTerraformError):
 
 class InvalidTeardownMode(InvalidOption):
     """Invalid Teardown Option Error"""
+
+
+class InvalidState(PytestTerraformError):
+    """Failure to load / parse state"""
+
+
+class ModuleNotFound(ValueError):
+    """module not found"""

pytest_terraform/tf.py

@@ -17,12 +17,13 @@
 import subprocess
 import sys
 from collections import UserString, defaultdict
+from typing import Any, Dict, Optional, Tuple, Union
 
 import jmespath
 import pytest
 from py.path import local
 
-from .exceptions import TerraformCommandFailed
+from .exceptions import InvalidState, ModuleNotFound, TerraformCommandFailed
 from .options import teardown as td
 
 
@@ -34,10 +35,6 @@ def find_binary(bin_name):
             return candidate
 
 
-class ModuleNotFound(ValueError):
-    """module not found"""
-
-
 class TerraformRunner(object):
 
     command_templates = {
@@ -83,7 +80,7 @@ def apply(self, plan=True):
         elif plan:
             apply_args = self._get_cmd_args("apply", plan="")
         self._run_cmd(apply_args)
-        return TerraformState.load(self.state_path)
+        return TerraformState.from_file(self.state_path)
 
     def plan(self, output=""):
         output = output and "-out=%s" % output or ""
@@ -118,26 +115,31 @@ def _run_cmd(self, args):
 
 class TerraformStateJson(UserString):
     @classmethod
-    def from_dict(cls, state):
+    def from_dict(cls, state: Dict[str, Any]):
+        """create TerraformStateJson from dictionary"""
         s = cls("")
         s.update_dict(state)
         return s
 
-    def update(self, state):
+    def update(self, state: str):
+        """update TerraformStateJson object with new data"""
         if not isinstance(state, str):
             raise ValueError(f"{state} is not a string")
 
         self.data = str(state)
 
-    def update_dict(self, state):
+    def update_dict(self, state: Dict[str, Any]):
+        """update TerraformStateJson from a dict"""
         self.update(json.dumps(state, indent=4))
 
     @property
     def dict(self):
+        """return the TerraformStateJson as a dict"""
         return json.loads(self.data)
 
     @dict.setter
-    def dict(self, data):
+    def dict(self, data: Dict[str, Any]):
+        """update TerraformStateJson from a dict"""
         try:
             self.update_dict(data)
         except (ValueError, TypeError):
@@ -195,20 +197,58 @@ def get(self, k, default=None):
         return default
 
     @classmethod
-    def load(cls, state):
-        resources = {}
-        outputs = {}
+    def from_file(cls, path: str):
+        """create TerraformState from a file
+
+        File can either be a Terraform Plan state, or a recorded
+        pytest-terraform state
+        """
+        if not os.path.isfile(path):
+            raise InvalidState("{} could not be located".format(path))
+
+        with open(path) as fh:
+            state = fh.read()
+
+        return cls.from_string(state)
+
+    @classmethod
+    def from_string(cls, state: Union[TerraformStateJson, str]):
+        """create TerraformState from string
+
+        State string can be a bytestring or a TerraformStateJson
+        string object
+        """
+        resources, outputs = cls.parse_state(state)
+        return cls(resources, outputs)
+
+    def update(self, state: Union[TerraformStateJson, str]):
+        """update TerraformState values"""
+        resources, outputs = self.parse_state(state)
+        self.resources = resources
+        self.outputs = outputs
 
+    @staticmethod
+    def parse_state(
+        state: Union[TerraformStateJson, str]
+    ) -> Tuple[Dict[str, any], Dict[str, Any]]:
+        """extract resources and outputs from state
+
+        where state is one of the following:
+        * Terraform state output as a string
+        * Recorded pytest-terraform state
+        * TerraformStateJson object
+        """
         if isinstance(state, TerraformStateJson):
             data = state.dict
-        elif os.path.isfile(state):
-            with open(state) as fh:
-                data = json.load(fh)
         else:
             data = json.loads(state)
 
         if "pytest-terraform" in data:
-            return cls(data["resources"], data["outputs"])
+            return (data["resources"], data["outputs"])
+
+        resources = {}
+        outputs = {}
+
         for r in data.get("resources", ()):
             rmap = resources.setdefault(r["type"], {})
             rmap[r["name"]] = dict(r["instances"][0]["attributes"])
@@ -225,9 +265,11 @@ def load(cls, state):
                     if "name" in kattr and vattr != rattrs["id"]:
                         rattrs[kattr] = vattr
                 rmap[rname] = rattrs
-        return cls(resources, outputs)
 
-    def save(self, state_path=None):
+        return (resources, outputs)
+
+    def save(self, state_path: Optional[str] = None) -> Optional[TerraformStateJson]:
+        """export state as a string or to a file"""
         state = {
             "pytest-terraform": 1,
             "outputs": self.outputs,
@@ -331,7 +373,9 @@ def __call__(self, request, tmpdir_factory, worker_id):
                 raise ValueError(
                     "Replay resources don't exist for %s" % self.tf_root_module
                 )
-            return TerraformTestApi.load(os.path.join(module_dir, "tf_resources.json"))
+            return TerraformTestApi.from_file(
+                os.path.join(module_dir, "tf_resources.json")
+            )
         work_dir = tmpdir_factory.mktemp(self.tf_root_module, numbered=True).join("work")
         self.runner = self.get_runner(module_dir, work_dir)
         return self.create(request, module_dir)
@@ -345,7 +389,7 @@ def create(self, request, module_dir):
             test_api = self.runner.apply()
             tfstatejson = test_api.save()
             self.config.hook.pytest_terraform_modify_state(tfstate=tfstatejson)
-            test_api.load(tfstatejson)
+            test_api.update(tfstatejson)
             test_api.save(module_dir.join("tf_resources.json"))
             return test_api
         except Exception:

tests/test_terraform.py

@@ -5,6 +5,7 @@
 
 import pytest
 from pytest_terraform import tf
+from pytest_terraform.exceptions import InvalidState
 
 
 def test_frame_walk():
@@ -42,7 +43,7 @@ def test_fixture_factory():
 
 
 def test_tf_resources():
-    state = tf.TerraformState.load(
+    state = tf.TerraformState.from_file(
         os.path.join(os.path.dirname(__file__), "burnify.tfstate")
     )
 
@@ -62,9 +63,9 @@ def test_tf_string_resources():
     with open(os.path.join(os.path.dirname(__file__), "burnify.tfstate")) as f:
         burnify = f.read()
 
-    state = tf.TerraformState.load(burnify)
+    state = tf.TerraformState.from_string(burnify)
     save_state = str(state.save())
-    reload = tf.TerraformState.load(save_state)
+    reload = tf.TerraformState.from_string(save_state)
 
     assert len(state.resources) == 9
     assert len(reload.resources) == 9
@@ -76,16 +77,21 @@ def test_tf_statejson_resources():
     with open(os.path.join(os.path.dirname(__file__), "burnify.tfstate")) as f:
         burnify = f.read()
 
-    state = tf.TerraformState.load(burnify)
+    state = tf.TerraformState.from_string(burnify)
     save_state = state.save()
-    reload = tf.TerraformState.load(save_state)
+    reload = tf.TerraformState.from_string(save_state)
 
     assert len(state.resources) == 9
     assert len(reload.resources) == 9
 
     assert save_state == reload.save()
 
 
+def test_tf_state_bad_file():
+    with pytest.raises(InvalidState):
+        tf.TerraformState.from_file("/not-exists1")
+
+
 def test_tf_statejson_from_dict():
     obj = {
         "test": "foo",