Add handling for initial/final states in "remove_node" graph function

resources/graphs/state_machine/graph.py

@@ -353,3 +353,22 @@ class StateMachineGraph(DirectedGraph):
         self._edges[connecting_edge.get_name()] = connecting_edge
 
         return connecting_edge
+
+    def remove_node(self, node_indentifier):
+        """
+        Removes node from graph.
+        :param node_indentifier: Node or name of node to remove.
+        :return: Removed node.
+        """
+        node = self.get_node(node_indentifier)
+
+        # Remove initial state, if present.
+        if node.is_initial:
+            self._initial_states.pop(node.get_name())
+
+        # Remove final state, if present.
+        if node.is_final:
+            self._final_states.pop(node.get_name())
+
+        # Call parent implementation.
+        return super().remove_node(node_indentifier)