Synnefo » snf-ganeti

import inspect

import re

import traceback

from cmdlib.testsupport.iallocator_mock import patchIAllocator

from ganeti import errors

from ganeti import opcodes

    * C{iallocator_cls}: @see L{patchIAllocator}

  REMOVE = object()

    self._iallocator_patcher = None

    self.ResetMocks()

  def _StopPatchers(self):

    if self._iallocator_patcher is not None:

      self._iallocator_patcher.stop()

      self._iallocator_patcher = None

  def tearDown(self):

    super(CmdlibTestCase, self).tearDown()

    self._StopPatchers()

  def _GetTestModule(self):

    module = inspect.getsourcefile(self.__class__).split("/")[-1]

    suffix = "_unittest.py"

    assert module.endswith(suffix), "Naming convention for cmdlib test" \

                                    " modules is: <module>%s (found '%s')"\

                                    % (suffix, module)

    return module[:-len(suffix)]

  def ResetMocks(self):

    """Resets all mocks back to their initial state.

    This is useful if you want to execute more than one opcode in a single

    test.

    """

    self._StopPatchers()

    try:

      self._iallocator_patcher = patchIAllocator(self._GetTestModule())

      self.iallocator_cls = self._iallocator_patcher.start()

    except ImportError:

      # this test module does not use iallocator, no patching performed

      self._iallocator_patcher = None

  def ExecOpCodeExpectException(self, opcode,

                                expected_exception,

                                expected_regex=None):

    """Executes the given opcode and expects an exception.

    @param opcode: @see L{ExecOpCode}

    @type expected_exception: class

    @param expected_exception: the exception which must be raised

    @type expected_regex: string

    @param expected_regex: if not C{None}, a regular expression which must be

          present in the string representation of the exception

    """

    try:

      self.ExecOpCode(opcode)

    except expected_exception, e:

      if expected_regex is not None:

        assert re.search(expected_regex, str(e)) is not None, \

                "Caught exception '%s' did not match '%s'" % \

                  (str(e), expected_regex)

    except Exception, e:

      tb = traceback.format_exc()

      raise AssertionError("%s\n(See original exception above)\n"

                           "Expected exception '%s' was not raised,"

                           " got '%s' of class '%s' instead." %

                           (tb, expected_exception, e, e.__class__))

    else:

      raise AssertionError("Expected exception '%s' was not raised" %

                           expected_exception)

  def ExecOpCodeExpectOpPrereqError(self, opcode, expected_regex=None):

    """Executes the given opcode and expects a L{errors.OpPrereqError}

    @see L{ExecOpCodeExpectException}

    """

    self.ExecOpCodeExpectException(opcode, errors.OpPrereqError, expected_regex)

  def ExecOpCodeExpectOpExecError(self, opcode, expected_regex=None):

    """Executes the given opcode and expects a L{errors.OpExecError}

    @see L{ExecOpCodeExpectException}

    """

    self.ExecOpCodeExpectException(opcode, errors.OpExecError, expected_regex)

  def CopyOpCode(self, opcode, **kwargs):

    """Creates a copy of the given opcode and applies modifications to it

    @type opcode: opcode.OpCode

    @param opcode: the opcode to copy

    @type kwargs: dict

    @param kwargs: dictionary of fields to overwrite in the copy. The special

          value L{REMOVE} can be used to remove fields from the copy.

    @return: a copy of the given opcode

    """

    state = opcode.__getstate__()

    for key, value in kwargs.items():

      if value == self.REMOVE:

        del state[key]

      else:

        state[key] = value

    return opcodes.OpCode.LoadOpCode(state)