Source code for utah.client.state_agent

# Ubuntu Testing Automation Harness
# Copyright 2012 Canonical Ltd.

# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License version 3, as published
# by the Free Software Foundation.

# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranties of
# MERCHANTABILITY, SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU General Public License for more details.

# You should have received a copy of the GNU General Public License along
# with this program.  If not, see <http://www.gnu.org/licenses/>.

"""Provide functionality for saving and restoring run state."""


import os

import yaml

from utah.client.common import DEFAULT_STATE_FILE, parse_yaml_file


[docs]class StateAgent(object):

    """State saving base class.

    Accepts a dictionary of state info and prints it to the state file.

    """

    state_file = DEFAULT_STATE_FILE

    def __init__(self, state_file=None):
        if state_file is not None:
            self.state_file = state_file

[docs]    def clean(self):
        """Clean up the state file if it exists."""
        # Don't fail if the state_file doesn't exists.
        try:
            os.remove(self.state_file)
        except OSError as e:
            if e.errno == 2:  # file not found
                pass

[docs]    def save_state(self, state):
        """Save state to the state_file."""
        with open(self.state_file, 'w') as fp:
            fp.write(str(state))

[docs]    def load_state(self):
        """Load state from the state_file.

        :returns: state information from file.
        :rtype: dict

        """
        state = {}

        if os.path.exists(self.state_file):
            with open(self.state_file, 'r') as fp:
                state = fp.read()

        return state


[docs]class StateAgentYAML(StateAgent):

    """YAML based state saver."""

[docs]    def save_state(self, state):
        """Output the state as YAML."""
        yaml_state = yaml.dump(state, default_flow_style=False)
        super(StateAgentYAML, self).save_state(yaml_state)

[docs]    def load_state(self):
        """Load state from YAML state_file.

        :returns: state information from YAML file
        :rtype: dict

        """
        state = {}

        if os.path.exists(self.state_file):
            state = parse_yaml_file(self.state_file)

        return state