Statistics
| Branch: | Tag: | Revision:

root / test / py / cmdlib / testsupport / rpc_runner_mock.py @ 4869595d

History | View | Annotate | Download (5.8 kB)

1
#
2
#
3

    
4
# Copyright (C) 2013 Google Inc.
5
#
6
# This program is free software; you can redistribute it and/or modify
7
# it under the terms of the GNU General Public License as published by
8
# the Free Software Foundation; either version 2 of the License, or
9
# (at your option) any later version.
10
#
11
# This program is distributed in the hope that it will be useful, but
12
# WITHOUT ANY WARRANTY; without even the implied warranty of
13
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14
# General Public License for more details.
15
#
16
# You should have received a copy of the GNU General Public License
17
# along with this program; if not, write to the Free Software
18
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
19
# 02110-1301, USA.
20

    
21

    
22
"""Support for mocking the RPC runner"""
23

    
24

    
25
import mock
26

    
27
from ganeti import objects
28
from ganeti.rpc import node as rpc
29

    
30
from cmdlib.testsupport.util import patchModule
31

    
32

    
33
def CreateRpcRunnerMock():
34
  """Creates a new L{mock.MagicMock} tailored for L{rpc.RpcRunner}
35

36
  """
37
  ret = mock.MagicMock(spec=rpc.RpcRunner)
38
  return ret
39

    
40

    
41
class RpcResultsBuilder(object):
42
  """Helper class which assists in constructing L{rpc.RpcResult} objects.
43

44
  This class provides some convenience methods for constructing L{rpc.RpcResult}
45
  objects. It is possible to create single results with the C{Create*} methods
46
  or to create multi-node results by repeatedly calling the C{Add*} methods and
47
  then obtaining the final result with C{Build}.
48

49
  The C{node} parameter of all the methods can either be a L{objects.Node}
50
  object, a node UUID or a node name. You have to provide the cluster config
51
  in the constructor if you want to use node UUID's/names.
52

53
  A typical usage of this class is as follows::
54

55
    self.rpc.call_some_rpc.return_value = \
56
      RpcResultsBuilder(cfg=self.cfg) \
57
        .AddSuccessfulNode(node1,
58
                           {
59
                             "result_key": "result_data",
60
                             "another_key": "other_data",
61
                           }) \
62
        .AddErrorNode(node2) \
63
        .Build()
64

65
  """
66

    
67
  def __init__(self, cfg=None, use_node_names=False):
68
    """Constructor.
69

70
    @type cfg: L{ganeti.config.ConfigWriter}
71
    @param cfg: used to resolve nodes if not C{None}
72
    @type use_node_names: bool
73
    @param use_node_names: if set to C{True}, the node field in the RPC results
74
          will contain the node name instead of the node UUID.
75
    """
76
    self._cfg = cfg
77
    self._use_node_names = use_node_names
78
    self._results = []
79

    
80
  def _GetNode(self, node_id):
81
    if isinstance(node_id, objects.Node):
82
      return node_id
83

    
84
    node = None
85
    if self._cfg is not None:
86
      node = self._cfg.GetNodeInfo(node_id)
87
      if node is None:
88
        node = self._cfg.GetNodeInfoByName(node_id)
89

    
90
    assert node is not None, "Failed to find '%s' in configuration" % node_id
91
    return node
92

    
93
  def _GetNodeId(self, node_id):
94
    node = self._GetNode(node_id)
95
    if self._use_node_names:
96
      return node.name
97
    else:
98
      return node.uuid
99

    
100
  def CreateSuccessfulNodeResult(self, node, data=None):
101
    """@see L{RpcResultsBuilder}
102

103
    @param node: @see L{RpcResultsBuilder}.
104
    @type data: dict
105
    @param data: the data as returned by the RPC
106
    @rtype: L{rpc.RpcResult}
107
    """
108
    if data is None:
109
      data = {}
110
    return rpc.RpcResult(data=(True, data), node=self._GetNodeId(node))
111

    
112
  def CreateFailedNodeResult(self, node):
113
    """@see L{RpcResultsBuilder}
114

115
    @param node: @see L{RpcResultsBuilder}.
116
    @rtype: L{rpc.RpcResult}
117
    """
118
    return rpc.RpcResult(failed=True, node=self._GetNodeId(node))
119

    
120
  def CreateOfflineNodeResult(self, node):
121
    """@see L{RpcResultsBuilder}
122

123
    @param node: @see L{RpcResultsBuilder}.
124
    @rtype: L{rpc.RpcResult}
125
    """
126
    return rpc.RpcResult(failed=True, offline=True, node=self._GetNodeId(node))
127

    
128
  def CreateErrorNodeResult(self, node, error_msg=None):
129
    """@see L{RpcResultsBuilder}
130

131
    @param node: @see L{RpcResultsBuilder}.
132
    @type error_msg: string
133
    @param error_msg: the error message as returned by the RPC
134
    @rtype: L{rpc.RpcResult}
135
    """
136
    return rpc.RpcResult(data=(False, error_msg), node=self._GetNodeId(node))
137

    
138
  def AddSuccessfulNode(self, node, data=None):
139
    """@see L{CreateSuccessfulNode}
140

141
    @rtype: L{RpcResultsBuilder}
142
    @return: self for chaining
143

144
    """
145
    self._results.append(self.CreateSuccessfulNodeResult(node, data))
146
    return self
147

    
148
  def AddFailedNode(self, node):
149
    """@see L{CreateFailedNode}
150

151
    @rtype: L{RpcResultsBuilder}
152
    @return: self for chaining
153

154
    """
155
    self._results.append(self.CreateFailedNodeResult(node))
156
    return self
157

    
158
  def AddOfflineNode(self, node):
159
    """@see L{CreateOfflineNode}
160

161
    @rtype: L{RpcResultsBuilder}
162
    @return: self for chaining
163

164
    """
165
    self._results.append(self.CreateOfflineNodeResult(node))
166
    return self
167

    
168
  def AddErrorNode(self, node, error_msg=None):
169
    """@see L{CreateErrorNode}
170

171
    @rtype: L{RpcResultsBuilder}
172
    @return: self for chaining
173

174
    """
175
    self._results.append(self.CreateErrorNodeResult(node, error_msg=error_msg))
176
    return self
177

    
178
  def Build(self):
179
    """Creates a dictionary holding multi-node results
180

181
    @rtype: dict
182
    """
183
    return dict((result.node, result) for result in self._results)
184

    
185

    
186
# pylint: disable=C0103
187
def patchRpc(module_under_test):
188
  """Patches the L{ganeti.rpc} module for tests.
189

190
  This function is meant to be used as a decorator for test methods.
191

192
  @type module_under_test: string
193
  @param module_under_test: the module within cmdlib which is tested. The
194
        "ganeti.cmdlib" prefix is optional.
195

196
  """
197
  return patchModule(module_under_test, "rpc", wraps=rpc)
198

    
199

    
200
def SetupDefaultRpcModuleMock(rpc_mod):
201
  """Configures the given rpc_mod.
202

203
  All relevant functions in rpc_mod are stubbed in a sensible way.
204

205
  @param rpc_mod: the mock module to configure
206

207
  """
208
  rpc_mod.DnsOnlyRunner.return_value = CreateRpcRunnerMock()