<commit_before>import numpy as np
import tensorflow as tf

# TODO: check the methods of _TensorLike
class TensorTrain(object):
  """Represents a Tensor Train object (a TT-tensor or TT-matrix).
  t3f represents a Tensor Train object as a tuple of TT-cores.
  ```
  @@__init__
  @@get_shape
  @@name
  @@tt_cores
  @@dtype
  @@op
  @@graph
  @@get_raw_shape
  @@ndims
  @@get_tt_ranks
  @@is_tt_matrix
  """

  def __init__(self, tt_cores, shape=None, tt_ranks=None, convert_to_tensors=True):
    """Creates a `TensorTrain`.
    Args:
      tt_cores: A tuple of 3d or 4d tensor-like objects of shape
        `[r_k-1, n_k, r_k]`.
        Tensor-like can be numpy array, tf.Tensor, of tf.Variable
      shape: Shape of the underlying tensor. If None, tries to infer from the
        cores (not always possible even if it should be, e.g. if ranks are
        unknown, than the whole shape of a core can be unknown).
      tt_ranks: a TensorShape of length d+1 (d is the dimensionality of
        the underlying tensor). The first and the last ranks are assumed to
        equal to 1. If None, tries to infer the ranks from the cores.
      convert_to_tensors: bool, if True than convert each element of the
        tt_cores tuple into a tf.Tensor (e.g. to initialize from np.array)

    Returns:
      A `TensorTrain`.

    Raises:
      ValueError if the provided TT-cores are not valid or inconsistent with
        the provided shape.
    """
    tt_cores = list(tt_cores)
    if convert_to_tensors:
      # TODO: what does this namescope do?
      with tf.name_scope("TensorTrain", tt_cores):
        # TODO: should we pass as_ref=True because we want to be able to update
        # values later if it is a VariableOp??
        for i in range(len(tt_cores)):
          name = "core%d" % i
          tt_cores[i] = tf.convert_to_tensor(
              tt_cores[i], name=name, as_ref=False)

    if not _are_tt_cores_valid(tt_cores, shape, tt_ranks):
      raise ValueError('the tt_cores provided to TensorTrain constructor are '
                       'not valid, have different dtypes, or are inconsistent '
                       'with the provided shape.')

    self._tt_cores = tuple(tt_cores)
    self._shape = _clean_shape(shape)
    self._tt_ranks = tt_ranks

  def get_raw_shape(self):
    """Get tuple of `TensorShapes` representing the shapes of the underlying TT-tensor.

    Tuple contains one `TensorShape` for TT-tensor and 2 `TensorShapes` for
    TT-matrix

    Returns:
      A tuple of `TensorShape` objects.
    """
    if self._shape is not None:
      return self._shape
    else:
      num_dims = self.ndims()
      num_tensor_shapes = len(self.tt_cores[0].get_shape().as_list()) - 2
      shapes = [[] for _ in range(num_tensor_shapes)]
      for dim in range(num_dims):
        curr_core_shape = self.tt_cores[dim].get_shape()
        for i in range(num_tensor_shapes):
          shapes[i].append(curr_core_shape[i + 1])
      for i in range(num_tensor_shapes):
        shapes[i] = tf.TensorShape(shapes[i])
      return tuple(shapes)

  def get_shape(self):
    """Get the `TensorShape` representing the shape of the dense tensor.
    Returns:
      A `TensorShape` object.
    """
    raw_shape = self.get_raw_shape()
    if self.is_tt_matrix():
      # TODO: return TensorShape.
      m = np.prod(raw_shape[0].as_list())
      n = np.prod(raw_shape[1].as_list())
      return m, n
    else:
      return self.get_raw_shape()[0]

  @property
  def tt_cores(self):
    """A tuple of TT-cores.
    Returns:
      A tuple of 3d or 4d tensors shape `[r_k-1, n_k, r_k]`.
    """
    return self._tt_cores

  @property
  def dtype(self):
    """The `DType` of elements in this tensor."""
    # TODO: where is this created?
    return self.tt_cores[0].dtype

  @property
  def name(self):
    """The name of the TensorTrain.

    Returns:
      String, the scope in which the TT-cores are defined.
    """
    core_name = self.tt_cores[0].name
    idx = core_name.rfind('/')
    return core_name[:idx]

  # TODO: it seems like instead of dense_shape we should use get_shape().
  # But maybe dense_shape() name is better?
  # @property
  # def dense_shape(self):
  #   """A 1-D Tensor of int64 representing the shape of the dense tensor."""
  #   return self._dense_shape

  @property
  def graph(self):
    """The `Graph` that contains the tt_cores tensors."""
    return self.tt_cores[0].graph

  def __str__(self):
    raise NotImplementedError
    # return "TensorTrain(indices=%s, values=%s, dense_shape=%s)" % (
    #     self._indices, self._values, self._dense_shape)

  def ndims(self):
    """Get the number of dimensions of the underlying TT-tensor.
    Returns:
      A number.
    """
    return len(self.tt_cores)

  def get_tt_ranks(self):
    """Get the TT-ranks in an array of size `num_dims`+1.

    The first and the last TT-rank are guarantied to be 1.

    Returns:
      TensorShape of size `num_dims`+1.
    """
    if self._tt_ranks is not None:
      return self._tt_ranks
    else:
      num_dims = self.ndims()
      ranks = []
      for i in range(num_dims):
        ranks.append(self.tt_cores[i].get_shape()[0])
      return tf.TensorShape(ranks)

  def is_tt_matrix(self):
    """Returns True if the TensorTrain object represents a TT-matrix.
    Returns:
      bool
    """
    if self._shape is not None:
      return len(self._shape) == 2
    return len(self.tt_cores[0].get_shape().as_list()) == 4

  def __getitem__(self, slice_spec):
    """Basic indexing, returns a `TensorTrain` containing the specified region.
    """
    new_tt_cores = []
    reminder = None
    for i in range(self.ndims()):
      curr_core = self.tt_cores[i]
      if self.is_tt_matrix():
        raise NotImplementedError
      else:
        sliced_core = curr_core[:, slice_spec[i], :]
        if len(curr_core.get_shape()) != len(sliced_core.get_shape()):
          # This index is specified exactly and we want to collapse this axis.
          if reminder is None:
            reminder = sliced_core
          else:
            reminder = tf.matmul(reminder, sliced_core)
        else:
          if reminder is not None:
            # Add reminder from the previous collapsed cores to the current
            # core.
            # TODO: is it bad to use as_list? E.g. if we dont now the ranks
            # on the graph construction stage.
            old_shape = sliced_core.get_shape().as_list()
            sliced_core = tf.reshape(sliced_core, (old_shape[0], -1))
            sliced_core = tf.matmul(reminder, sliced_core)
            sliced_core = tf.reshape(sliced_core, (-1, old_shape[1], old_shape[2]))
            reminder = None
          new_tt_cores.append(sliced_core)

    if reminder is not None:
      # The reminder obtained from collapsing the last cores.
      old_shape = new_tt_cores[-1].get_shape().as_list()
      new_tt_cores[-1] = tf.reshape(new_tt_cores[-1], (-1, old_shape[-1]))
      new_tt_cores[-1] = tf.matmul(new_tt_cores[-1], reminder)
      new_tt_cores[-1] = tf.reshape(new_tt_cores[-1], (old_shape[0], old_shape[1], 1))
      reminder = None
    return TensorTrain(new_tt_cores)

  def eval(self, feed_dict=None, session=None):
    """Evaluates this sparse tensor in a `Session`.
    Calling this method will execute all preceding operations that
    produce the inputs needed for the operation that produces this
    tensor.
    *N.B.* Before invoking `SparseTensor.eval()`, its graph must have been
    launched in a session, and either a default session must be
    available, or `session` must be specified explicitly.
    Args:
      feed_dict: A dictionary that maps `Tensor` objects to feed values.
        See [`Session.run()`](../../api_docs/python/client.md#Session.run) for a
        description of the valid feed values.
      session: (Optional.) The `Session` to be used to evaluate this sparse
        tensor. If none, the default session will be used.
    Returns:
      ????
    """
    # TODO: what to return, None?
    if session is None:
      session = tf.get_default_session()
    session.run(self.tt_cores)

  # TODO: do we need this?
  # @staticmethod
  # def _override_operator(operator, func):
  #   _override_helper(SparseTensor, operator, func)

def _clean_shape(shape):
  """Returns a tuple of TensorShapes for any valid shape representation.

  Args:
    shape: An np.array, a tf.TensorShape (for tensors), a tuple of
      tf.TensorShapes (for TT-matrices or tensors), or None

  Returns:
    A tuple of tf.TensorShape, or None is the input is None
  """
  if shape is None:
    return None
  if isinstance(shape, tf.TensorShape) or isinstance(shape[0], tf.TensorShape):
    # Assume tf.TensorShape.
    if isinstance(shape, tf.TensorShape):
      shape = tuple((shape,))
  else:
    shape = np.array(shape)
    # Make sure that the shape is 2-d array both for tensors and TT-matrices.
    shape = np.squeeze(shape)
    if len(shape.shape) == 1:
      # A tensor.
      shape = [shape]
    for i in range(len(shape)):
      shape[i] = tf.TensorShape(shape[i])
    shape = tuple(shape)
  return shape

def _are_tt_cores_valid(tt_cores, shape, tt_ranks):
  """Check if dimensions of the TT-cores are consistent and the dtypes coincide.

  Args:
    tt_cores: a tuple of `Tensor` objects
    shape: An np.array, a tf.TensorShape (for tensors), a tuple of
      tf.TensorShapes (for TT-matrices or tensors), or None
    tt_ranks: An np.array or a tf.TensorShape of length len(tt_cores)+1.

  Returns:
    boolean, True if the dimensions and dtypes are consistent.
  """
  shape = _clean_shape(shape)
  num_dims = len(tt_cores)

  for core_idx in range(1, num_dims):
    if tt_cores[core_idx].dtype != tt_cores[0].dtype:
      return False
  try:
    for core_idx in range(num_dims):
      curr_core_shape = tt_cores[core_idx].get_shape()
      if len(curr_core_shape) != len(tt_cores[0].get_shape()):
        # Shapes are inconsistent.
        return False
      if shape is not None:
        for i in range(len(shape)):
          if curr_core_shape[i + 1] != shape[i][core_idx]:
            # The TT-cores are not aligned with the given shape.
            return False
      if core_idx >= 1:
        prev_core_shape = tt_cores[core_idx - 1].get_shape()
        if curr_core_shape[0] != prev_core_shape[-1]:
          # TT-ranks are inconsistent.
          return False
      # print(core_idx)
      if tt_ranks is not None:
        if curr_core_shape[0] != tt_ranks[core_idx]:
          # The TT-ranks are not aligned with the TT-cores shape.
          return False
        if curr_core_shape[-1] != tt_ranks[core_idx + 1]:
          # The TT-ranks are not aligned with the TT-cores shape.
          return False
    if tt_cores[0].get_shape()[0] != 1 or tt_cores[-1].get_shape()[-1] != 1:
      # The first or the last rank is not 1.
      return False
  except ValueError:
    # The shape of the TT-cores is undetermined, can not validate it.
    pass
  return True
<commit_msg>Fix bug with assigning TensorShape object to elements of numeric np.array<commit_after>import numpy as np
import tensorflow as tf

# TODO: check the methods of _TensorLike
class TensorTrain(object):
  """Represents a Tensor Train object (a TT-tensor or TT-matrix).
  t3f represents a Tensor Train object as a tuple of TT-cores.
  ```
  @@__init__
  @@get_shape
  @@name
  @@tt_cores
  @@dtype
  @@op
  @@graph
  @@get_raw_shape
  @@ndims
  @@get_tt_ranks
  @@is_tt_matrix
  """

  def __init__(self, tt_cores, shape=None, tt_ranks=None, convert_to_tensors=True):
    """Creates a `TensorTrain`.
    Args:
      tt_cores: A tuple of 3d or 4d tensor-like objects of shape
        `[r_k-1, n_k, r_k]`.
        Tensor-like can be numpy array, tf.Tensor, of tf.Variable
      shape: Shape of the underlying tensor. If None, tries to infer from the
        cores (not always possible even if it should be, e.g. if ranks are
        unknown, than the whole shape of a core can be unknown).
      tt_ranks: a TensorShape of length d+1 (d is the dimensionality of
        the underlying tensor). The first and the last ranks are assumed to
        equal to 1. If None, tries to infer the ranks from the cores.
      convert_to_tensors: bool, if True than convert each element of the
        tt_cores tuple into a tf.Tensor (e.g. to initialize from np.array)

    Returns:
      A `TensorTrain`.

    Raises:
      ValueError if the provided TT-cores are not valid or inconsistent with
        the provided shape.
    """
    tt_cores = list(tt_cores)
    if convert_to_tensors:
      # TODO: what does this namescope do?
      with tf.name_scope("TensorTrain", tt_cores):
        # TODO: should we pass as_ref=True because we want to be able to update
        # values later if it is a VariableOp??
        for i in range(len(tt_cores)):
          name = "core%d" % i
          tt_cores[i] = tf.convert_to_tensor(
              tt_cores[i], name=name, as_ref=False)

    if not _are_tt_cores_valid(tt_cores, shape, tt_ranks):
      raise ValueError('the tt_cores provided to TensorTrain constructor are '
                       'not valid, have different dtypes, or are inconsistent '
                       'with the provided shape.')

    self._tt_cores = tuple(tt_cores)
    self._shape = _clean_shape(shape)
    self._tt_ranks = tt_ranks

  def get_raw_shape(self):
    """Get tuple of `TensorShapes` representing the shapes of the underlying TT-tensor.

    Tuple contains one `TensorShape` for TT-tensor and 2 `TensorShapes` for
    TT-matrix

    Returns:
      A tuple of `TensorShape` objects.
    """
    if self._shape is not None:
      return self._shape
    else:
      num_dims = self.ndims()
      num_tensor_shapes = len(self.tt_cores[0].get_shape().as_list()) - 2
      shapes = [[] for _ in range(num_tensor_shapes)]
      for dim in range(num_dims):
        curr_core_shape = self.tt_cores[dim].get_shape()
        for i in range(num_tensor_shapes):
          shapes[i].append(curr_core_shape[i + 1])
      for i in range(num_tensor_shapes):
        shapes[i] = tf.TensorShape(shapes[i])
      return tuple(shapes)

  def get_shape(self):
    """Get the `TensorShape` representing the shape of the dense tensor.
    Returns:
      A `TensorShape` object.
    """
    raw_shape = self.get_raw_shape()
    if self.is_tt_matrix():
      # TODO: return TensorShape.
      m = np.prod(raw_shape[0].as_list())
      n = np.prod(raw_shape[1].as_list())
      return m, n
    else:
      return self.get_raw_shape()[0]

  @property
  def tt_cores(self):
    """A tuple of TT-cores.
    Returns:
      A tuple of 3d or 4d tensors shape `[r_k-1, n_k, r_k]`.
    """
    return self._tt_cores

  @property
  def dtype(self):
    """The `DType` of elements in this tensor."""
    # TODO: where is this created?
    return self.tt_cores[0].dtype

  @property
  def name(self):
    """The name of the TensorTrain.

    Returns:
      String, the scope in which the TT-cores are defined.
    """
    core_name = self.tt_cores[0].name
    idx = core_name.rfind('/')
    return core_name[:idx]

  # TODO: it seems like instead of dense_shape we should use get_shape().
  # But maybe dense_shape() name is better?
  # @property
  # def dense_shape(self):
  #   """A 1-D Tensor of int64 representing the shape of the dense tensor."""
  #   return self._dense_shape

  @property
  def graph(self):
    """The `Graph` that contains the tt_cores tensors."""
    return self.tt_cores[0].graph

  def __str__(self):
    raise NotImplementedError
    # return "TensorTrain(indices=%s, values=%s, dense_shape=%s)" % (
    #     self._indices, self._values, self._dense_shape)

  def ndims(self):
    """Get the number of dimensions of the underlying TT-tensor.
    Returns:
      A number.
    """
    return len(self.tt_cores)

  def get_tt_ranks(self):
    """Get the TT-ranks in an array of size `num_dims`+1.

    The first and the last TT-rank are guarantied to be 1.

    Returns:
      TensorShape of size `num_dims`+1.
    """
    if self._tt_ranks is not None:
      return self._tt_ranks
    else:
      num_dims = self.ndims()
      ranks = []
      for i in range(num_dims):
        ranks.append(self.tt_cores[i].get_shape()[0])
      return tf.TensorShape(ranks)

  def is_tt_matrix(self):
    """Returns True if the TensorTrain object represents a TT-matrix.
    Returns:
      bool
    """
    if self._shape is not None:
      return len(self._shape) == 2
    return len(self.tt_cores[0].get_shape().as_list()) == 4

  def __getitem__(self, slice_spec):
    """Basic indexing, returns a `TensorTrain` containing the specified region.
    """
    new_tt_cores = []
    reminder = None
    for i in range(self.ndims()):
      curr_core = self.tt_cores[i]
      if self.is_tt_matrix():
        raise NotImplementedError
      else:
        sliced_core = curr_core[:, slice_spec[i], :]
        if len(curr_core.get_shape()) != len(sliced_core.get_shape()):
          # This index is specified exactly and we want to collapse this axis.
          if reminder is None:
            reminder = sliced_core
          else:
            reminder = tf.matmul(reminder, sliced_core)
        else:
          if reminder is not None:
            # Add reminder from the previous collapsed cores to the current
            # core.
            # TODO: is it bad to use as_list? E.g. if we dont now the ranks
            # on the graph construction stage.
            old_shape = sliced_core.get_shape().as_list()
            sliced_core = tf.reshape(sliced_core, (old_shape[0], -1))
            sliced_core = tf.matmul(reminder, sliced_core)
            sliced_core = tf.reshape(sliced_core, (-1, old_shape[1], old_shape[2]))
            reminder = None
          new_tt_cores.append(sliced_core)

    if reminder is not None:
      # The reminder obtained from collapsing the last cores.
      old_shape = new_tt_cores[-1].get_shape().as_list()
      new_tt_cores[-1] = tf.reshape(new_tt_cores[-1], (-1, old_shape[-1]))
      new_tt_cores[-1] = tf.matmul(new_tt_cores[-1], reminder)
      new_tt_cores[-1] = tf.reshape(new_tt_cores[-1], (old_shape[0], old_shape[1], 1))
      reminder = None
    return TensorTrain(new_tt_cores)

  def eval(self, feed_dict=None, session=None):
    """Evaluates this sparse tensor in a `Session`.
    Calling this method will execute all preceding operations that
    produce the inputs needed for the operation that produces this
    tensor.
    *N.B.* Before invoking `SparseTensor.eval()`, its graph must have been
    launched in a session, and either a default session must be
    available, or `session` must be specified explicitly.
    Args:
      feed_dict: A dictionary that maps `Tensor` objects to feed values.
        See [`Session.run()`](../../api_docs/python/client.md#Session.run) for a
        description of the valid feed values.
      session: (Optional.) The `Session` to be used to evaluate this sparse
        tensor. If none, the default session will be used.
    Returns:
      ????
    """
    # TODO: what to return, None?
    if session is None:
      session = tf.get_default_session()
    session.run(self.tt_cores)

  # TODO: do we need this?
  # @staticmethod
  # def _override_operator(operator, func):
  #   _override_helper(SparseTensor, operator, func)

def _clean_shape(shape):
  """Returns a tuple of TensorShapes for any valid shape representation.

  Args:
    shape: An np.array, a tf.TensorShape (for tensors), a tuple of
      tf.TensorShapes (for TT-matrices or tensors), or None

  Returns:
    A tuple of tf.TensorShape, or None is the input is None
  """
  if shape is None:
    return None
  if isinstance(shape, tf.TensorShape) or isinstance(shape[0], tf.TensorShape):
    # Assume tf.TensorShape.
    if isinstance(shape, tf.TensorShape):
      shape = tuple((shape,))
  else:
    np_shape = np.array(shape)
    # Make sure that the shape is 2-d array both for tensors and TT-matrices.
    np_shape = np.squeeze(np_shape)
    if len(np_shape.shape) == 1:
      # A tensor.
      np_shape = [np_shape]
    shape = []
    for i in range(len(np_shape)):
      shape.append(tf.TensorShape(np_shape[i]))
    shape = tuple(shape)
  return shape

def _are_tt_cores_valid(tt_cores, shape, tt_ranks):
  """Check if dimensions of the TT-cores are consistent and the dtypes coincide.

  Args:
    tt_cores: a tuple of `Tensor` objects
    shape: An np.array, a tf.TensorShape (for tensors), a tuple of
      tf.TensorShapes (for TT-matrices or tensors), or None
    tt_ranks: An np.array or a tf.TensorShape of length len(tt_cores)+1.

  Returns:
    boolean, True if the dimensions and dtypes are consistent.
  """
  shape = _clean_shape(shape)
  num_dims = len(tt_cores)

  for core_idx in range(1, num_dims):
    if tt_cores[core_idx].dtype != tt_cores[0].dtype:
      return False
  try:
    for core_idx in range(num_dims):
      curr_core_shape = tt_cores[core_idx].get_shape()
      if len(curr_core_shape) != len(tt_cores[0].get_shape()):
        # Shapes are inconsistent.
        return False
      if shape is not None:
        for i in range(len(shape)):
          if curr_core_shape[i + 1] != shape[i][core_idx]:
            # The TT-cores are not aligned with the given shape.
            return False
      if core_idx >= 1:
        prev_core_shape = tt_cores[core_idx - 1].get_shape()
        if curr_core_shape[0] != prev_core_shape[-1]:
          # TT-ranks are inconsistent.
          return False
      # print(core_idx)
      if tt_ranks is not None:
        if curr_core_shape[0] != tt_ranks[core_idx]:
          # The TT-ranks are not aligned with the TT-cores shape.
          return False
        if curr_core_shape[-1] != tt_ranks[core_idx + 1]:
          # The TT-ranks are not aligned with the TT-cores shape.
          return False
    if tt_cores[0].get_shape()[0] != 1 or tt_cores[-1].get_shape()[-1] != 1:
      # The first or the last rank is not 1.
      return False
  except ValueError:
    # The shape of the TT-cores is undetermined, can not validate it.
    pass
  return True
<|endoftext|>


<commit_before># Standard python modules.
import time
import sys

# citest modules.
import citest.gcp_testing as gcp
import citest.json_contract as jc
import citest.service_testing as st

# Spinnaker modules.
import spinnaker_testing as sk
import spinnaker_testing.gate as gate

class ServerGroupTestScenario(sk.SpinnakerTestScenario):

  @classmethod
  def new_agent(cls, bindings):
    '''Implements the base class interface to create a new agent.

    This method is called by the base classes during setup/initialization.

    Args:
      bindings: The bindings dictionary with configuration information
        that this factory can draw from to initialize. If the factory would
        like additional custom bindings it could add them to initArgumentParser.

    Returns:
      A citest.service_testing.BaseAgent that can interact with Gate.
      This is the agent that test operations will be posted to.
    '''
    return gate.new_agent(bindings)

  def __init__(self, bindings, agent=None):
    super(ServerGroupTestScenario, self).__init__(bindings, agent)

    # Our application name and path to post events to.
    self.TEST_APP = bindings['TEST_APP']
    self.__path = 'applications/%s/tasks' % self.TEST_APP

    # The spinnaker stack decorator for our resources.
    self.TEST_STACK = bindings['TEST_STACK']

    self.TEST_REGION = bindings['TEST_GCE_REGION']
    self.TEST_ZONE = bindings['TEST_GCE_ZONE']

    # Resource names used among tests.
    self.__cluster_name = '%s-%s' % (self.TEST_APP, self.TEST_STACK)
    self.__server_group_name = '%s-v000' % self.__cluster_name
    self.__cloned_server_group_name = '%s-v001' % self.__cluster_name
    self.__lb_name = '%s-%s-fe' % (self.TEST_APP, self.TEST_STACK)

  def create_load_balancer(self):
    job = [{
      'cloudProvider': 'gce',
      'loadBalancerName': self.__lb_name,
      'ipProtocol': 'TCP',
      'portRange': '8080',
      'provider': 'gce',
      'stack': self.TEST_STACK,
      'detail': 'frontend',
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'region': self.TEST_REGION,
      'listeners': [{
        'protocol': 'TCP',
        'portRange': '8080',
        'healthCheck': False
      }],
      'name': self.__lb_name,
      'type': 'upsertLoadBalancer',
      'availabilityZones': {self.TEST_REGION: []},
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Load Balancer Created', retryable_for_secs=30)
     .list_resources('forwarding-rules')
     .contains_path_value('name', self.__lb_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - create load balancer',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='create_load_balancer', data=payload, path=self.__path),
      contract=builder.build())

  def create_instances(self):
    job = [{
      'application': self.TEST_APP,
      'stack': self.TEST_STACK,
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'zone': self.TEST_ZONE,
      'network': 'default',
      'targetSize': 1,
      'capacity': {
        'min': 1,
        'max': 1,
        'desired': 1
      },
      'availabilityZones': {
        self.TEST_REGION: [self.TEST_ZONE]
      },
      'loadBalancers': [self.__lb_name],
      'instanceMetadata': {
        'load-balancer-names': self.__lb_name
      },
      'cloudProvider': 'gce',
      'image': self.bindings['TEST_GCE_IMAGE_NAME'],
      'instanceType': 'f1-micro',
      'initialNumReplicas': 1,
      'type': 'createServerGroup',
      'account': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Instance Created', retryable_for_secs=150)
     .list_resources('instance-groups')
     .contains_path_value('name', self.__server_group_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job,
        description='Server Group Tests - create initial server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='create_instances', data=payload, path=self.__path),
      contract=builder.build())

  def resize_server_group(self):
    job = [{
      'targetSize': 2,
      'capacity': {
        'min': 2,
        'max': 2,
        'desired': 2
      },
      'replicaPoolName': self.__server_group_name,
      'numReplicas': 2,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'asgName': self.__server_group_name,
      'type': 'resizeServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'cloudProvider': 'gce',
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Resized', retryable_for_secs=90)
     .inspect_resource('instance-groups',
                       self.__server_group_name,
                       ['--zone', self.TEST_ZONE])
     .contains_path_eq('size', 2))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - resize to 2 instances',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='resize_instances', data=payload, path=self.__path),
      contract=builder.build())

  def clone_server_group(self):
    job = [{
      'application': self.TEST_APP,
      'stack': self.TEST_STACK,
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'loadBalancers': [self.__lb_name],
      'targetSize': 1,
      'capacity': {
        'min': 1,
        'max': 1,
        'desired': 1
      },
      'zone': self.TEST_ZONE,
      'network': 'default',
      'instanceMetadata': {'load-balancer-names': self.__lb_name},
      'availabilityZones': {self.TEST_REGION: [self.TEST_ZONE]},
      'cloudProvider': 'gce',
      'source': {
        'account': self.bindings['GCE_CREDENTIALS'],
        'region': self.TEST_REGION,
        'zone': self.TEST_ZONE,
        'serverGroupName': self.__server_group_name,
        'asgName': self.__server_group_name
      },
      'instanceType': 'f1-micro',
      'image': self.bindings['TEST_GCE_IMAGE_NAME'],
      'initialNumReplicas': 1,
      'loadBalancers': [self.__lb_name],
      'type': 'cloneServerGroup',
      'account': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Cloned', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_path_value('baseInstanceName', self.__cloned_server_group_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - clone server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='clone_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def disable_server_group(self):
    job = [{
      'cloudProvider': 'gce',
      'asgName': self.__server_group_name,
      'serverGroupName': self.__server_group_name,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'disableServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Disabled', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_path_value('baseInstanceName', self.__server_group_name)
     .excludes_pred_list([
         jc.PathContainsPredicate('baseInstanceName', self.__server_group_name),
         jc.PathContainsPredicate('targetPools', 'https')]))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - disable server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='disable_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def enable_server_group(self):
    job = [{
      'cloudProvider': 'gce',
      'asgName': self.__server_group_name,
      'serverGroupName': self.__server_group_name,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'enableServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Enabled', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_pred_list([
         jc.PathContainsPredicate('baseInstanceName', self.__server_group_name),
         jc.PathContainsPredicate('targetPools', 'https')]))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - enable server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='enable_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def destroy_server_group(self, version):
    serverGroupName = '%s-%s' % (self.__cluster_name, version)
    job = [{
      'cloudProvider': 'gce',
      'asgName': serverGroupName,
      'serverGroupName': serverGroupName,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'destroyServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Destroyed', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .excludes_path_value('baseInstanceName', serverGroupName))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - destroy server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='destroy_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def delete_load_balancer(self):
    job = [{
      "loadBalancerName": self.__lb_name,
      "networkLoadBalancerName": self.__lb_name,
      "region": "us-central1",
      "type": "deleteLoadBalancer",
      "regions": ["us-central1"],
      "credentials": self.bindings['GCE_CREDENTIALS'],
      "cloudProvider": "gce",
      "user": "integration-tests"
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Load Balancer Created', retryable_for_secs=30)
     .list_resources('forwarding-rules')
     .excludes_path_value('name', self.__lb_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - delete load balancer',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='delete_load_balancer', data=payload, path=self.__path),
      contract=builder.build())

class ServerGroupTest(st.AgentTestCase):
  def test_a_create_load_balancer(self):
    self.run_test_case(self.scenario.create_load_balancer())

  def test_b_create_server_group(self):
    self.run_test_case(self.scenario.create_instances())

  def test_c_resize_server_group(self):
    self.run_test_case(self.scenario.resize_server_group())

  def test_d_clone_server_group(self):
    self.run_test_case(self.scenario.clone_server_group())

  def test_e_disable_server_group(self):
    self.run_test_case(self.scenario.disable_server_group())

  def test_f_enable_server_group(self):
    self.run_test_case(self.scenario.enable_server_group())

  def test_g_destroy_server_group_v000(self):
    self.run_test_case(self.scenario.destroy_server_group('v000'))

  def test_h_destroy_server_group_v001(self):
    self.run_test_case(self.scenario.destroy_server_group('v001'))

  def test_z_delete_load_balancer(self):
    self.run_test_case(self.scenario.delete_load_balancer())

def main():

  defaults = {
    'TEST_STACK': 'svrgrptst' + ServerGroupTestScenario.DEFAULT_TEST_ID,
    'TEST_APP': 'servergrouptest' + ServerGroupTestScenario.DEFAULT_TEST_ID
  }

  return st.ScenarioTestRunner.main(
      ServerGroupTestScenario,
      default_binding_overrides=defaults,
      test_case_list=[ServerGroupTest])

if __name__ == '__main__':
  sys.exit(main())
<commit_msg>Retry the clone_server_group test to mitigate flakiness.<commit_after># Standard python modules.
import time
import sys

# citest modules.
import citest.gcp_testing as gcp
import citest.json_contract as jc
import citest.service_testing as st

# Spinnaker modules.
import spinnaker_testing as sk
import spinnaker_testing.gate as gate

class ServerGroupTestScenario(sk.SpinnakerTestScenario):

  @classmethod
  def new_agent(cls, bindings):
    '''Implements the base class interface to create a new agent.

    This method is called by the base classes during setup/initialization.

    Args:
      bindings: The bindings dictionary with configuration information
        that this factory can draw from to initialize. If the factory would
        like additional custom bindings it could add them to initArgumentParser.

    Returns:
      A citest.service_testing.BaseAgent that can interact with Gate.
      This is the agent that test operations will be posted to.
    '''
    return gate.new_agent(bindings)

  def __init__(self, bindings, agent=None):
    super(ServerGroupTestScenario, self).__init__(bindings, agent)

    # Our application name and path to post events to.
    self.TEST_APP = bindings['TEST_APP']
    self.__path = 'applications/%s/tasks' % self.TEST_APP

    # The spinnaker stack decorator for our resources.
    self.TEST_STACK = bindings['TEST_STACK']

    self.TEST_REGION = bindings['TEST_GCE_REGION']
    self.TEST_ZONE = bindings['TEST_GCE_ZONE']

    # Resource names used among tests.
    self.__cluster_name = '%s-%s' % (self.TEST_APP, self.TEST_STACK)
    self.__server_group_name = '%s-v000' % self.__cluster_name
    self.__cloned_server_group_name = '%s-v001' % self.__cluster_name
    self.__lb_name = '%s-%s-fe' % (self.TEST_APP, self.TEST_STACK)

  def create_load_balancer(self):
    job = [{
      'cloudProvider': 'gce',
      'loadBalancerName': self.__lb_name,
      'ipProtocol': 'TCP',
      'portRange': '8080',
      'provider': 'gce',
      'stack': self.TEST_STACK,
      'detail': 'frontend',
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'region': self.TEST_REGION,
      'listeners': [{
        'protocol': 'TCP',
        'portRange': '8080',
        'healthCheck': False
      }],
      'name': self.__lb_name,
      'type': 'upsertLoadBalancer',
      'availabilityZones': {self.TEST_REGION: []},
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Load Balancer Created', retryable_for_secs=30)
     .list_resources('forwarding-rules')
     .contains_path_value('name', self.__lb_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - create load balancer',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='create_load_balancer', data=payload, path=self.__path),
      contract=builder.build())

  def create_instances(self):
    job = [{
      'application': self.TEST_APP,
      'stack': self.TEST_STACK,
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'zone': self.TEST_ZONE,
      'network': 'default',
      'targetSize': 1,
      'capacity': {
        'min': 1,
        'max': 1,
        'desired': 1
      },
      'availabilityZones': {
        self.TEST_REGION: [self.TEST_ZONE]
      },
      'loadBalancers': [self.__lb_name],
      'instanceMetadata': {
        'load-balancer-names': self.__lb_name
      },
      'cloudProvider': 'gce',
      'image': self.bindings['TEST_GCE_IMAGE_NAME'],
      'instanceType': 'f1-micro',
      'initialNumReplicas': 1,
      'type': 'createServerGroup',
      'account': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Instance Created', retryable_for_secs=150)
     .list_resources('instance-groups')
     .contains_path_value('name', self.__server_group_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job,
        description='Server Group Tests - create initial server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='create_instances', data=payload, path=self.__path),
      contract=builder.build())

  def resize_server_group(self):
    job = [{
      'targetSize': 2,
      'capacity': {
        'min': 2,
        'max': 2,
        'desired': 2
      },
      'replicaPoolName': self.__server_group_name,
      'numReplicas': 2,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'asgName': self.__server_group_name,
      'type': 'resizeServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'cloudProvider': 'gce',
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Resized', retryable_for_secs=90)
     .inspect_resource('instance-groups',
                       self.__server_group_name,
                       ['--zone', self.TEST_ZONE])
     .contains_path_eq('size', 2))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - resize to 2 instances',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='resize_instances', data=payload, path=self.__path),
      contract=builder.build())

  def clone_server_group(self):
    job = [{
      'application': self.TEST_APP,
      'stack': self.TEST_STACK,
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'loadBalancers': [self.__lb_name],
      'targetSize': 1,
      'capacity': {
        'min': 1,
        'max': 1,
        'desired': 1
      },
      'zone': self.TEST_ZONE,
      'network': 'default',
      'instanceMetadata': {'load-balancer-names': self.__lb_name},
      'availabilityZones': {self.TEST_REGION: [self.TEST_ZONE]},
      'cloudProvider': 'gce',
      'source': {
        'account': self.bindings['GCE_CREDENTIALS'],
        'region': self.TEST_REGION,
        'zone': self.TEST_ZONE,
        'serverGroupName': self.__server_group_name,
        'asgName': self.__server_group_name
      },
      'instanceType': 'f1-micro',
      'image': self.bindings['TEST_GCE_IMAGE_NAME'],
      'initialNumReplicas': 1,
      'loadBalancers': [self.__lb_name],
      'type': 'cloneServerGroup',
      'account': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Cloned', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_path_value('baseInstanceName', self.__cloned_server_group_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - clone server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='clone_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def disable_server_group(self):
    job = [{
      'cloudProvider': 'gce',
      'asgName': self.__server_group_name,
      'serverGroupName': self.__server_group_name,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'disableServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Disabled', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_path_value('baseInstanceName', self.__server_group_name)
     .excludes_pred_list([
         jc.PathContainsPredicate('baseInstanceName', self.__server_group_name),
         jc.PathContainsPredicate('targetPools', 'https')]))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - disable server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='disable_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def enable_server_group(self):
    job = [{
      'cloudProvider': 'gce',
      'asgName': self.__server_group_name,
      'serverGroupName': self.__server_group_name,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'enableServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Enabled', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .contains_pred_list([
         jc.PathContainsPredicate('baseInstanceName', self.__server_group_name),
         jc.PathContainsPredicate('targetPools', 'https')]))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - enable server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='enable_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def destroy_server_group(self, version):
    serverGroupName = '%s-%s' % (self.__cluster_name, version)
    job = [{
      'cloudProvider': 'gce',
      'asgName': serverGroupName,
      'serverGroupName': serverGroupName,
      'region': self.TEST_REGION,
      'zone': self.TEST_ZONE,
      'type': 'destroyServerGroup',
      'regions': [self.TEST_REGION],
      'zones': [self.TEST_ZONE],
      'credentials': self.bindings['GCE_CREDENTIALS'],
      'user': 'integration-tests'
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Server Group Destroyed', retryable_for_secs=90)
     .list_resources('managed-instance-groups')
     .excludes_path_value('baseInstanceName', serverGroupName))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - destroy server group',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='destroy_server_group', data=payload, path=self.__path),
      contract=builder.build())

  def delete_load_balancer(self):
    job = [{
      "loadBalancerName": self.__lb_name,
      "networkLoadBalancerName": self.__lb_name,
      "region": "us-central1",
      "type": "deleteLoadBalancer",
      "regions": ["us-central1"],
      "credentials": self.bindings['GCE_CREDENTIALS'],
      "cloudProvider": "gce",
      "user": "integration-tests"
    }]

    builder = gcp.GceContractBuilder(self.gce_observer)
    (builder.new_clause_builder('Load Balancer Created', retryable_for_secs=30)
     .list_resources('forwarding-rules')
     .excludes_path_value('name', self.__lb_name))

    payload = self.agent.make_json_payload_from_kwargs(
        job=job, description='Server Group Tests - delete load balancer',
        application=self.TEST_APP)

    return st.OperationContract(
      self.new_post_operation(
          title='delete_load_balancer', data=payload, path=self.__path),
      contract=builder.build())

class ServerGroupTest(st.AgentTestCase):
  def test_a_create_load_balancer(self):
    self.run_test_case(self.scenario.create_load_balancer())

  def test_b_create_server_group(self):
    self.run_test_case(self.scenario.create_instances())

  def test_c_resize_server_group(self):
    self.run_test_case(self.scenario.resize_server_group())

  def test_d_clone_server_group(self):
    self.run_test_case(self.scenario.clone_server_group(),
                       # TODO(ewiseblatt): 20160314
                       # There is a lock contention race condition
                       # in clouddriver that causes intermittent failure.
                       max_retries=5)

  def test_e_disable_server_group(self):
    self.run_test_case(self.scenario.disable_server_group())

  def test_f_enable_server_group(self):
    self.run_test_case(self.scenario.enable_server_group())

  def test_g_destroy_server_group_v000(self):
    self.run_test_case(self.scenario.destroy_server_group('v000'))

  def test_h_destroy_server_group_v001(self):
    self.run_test_case(self.scenario.destroy_server_group('v001'))

  def test_z_delete_load_balancer(self):
    self.run_test_case(self.scenario.delete_load_balancer())

def main():

  defaults = {
    'TEST_STACK': 'svrgrptst' + ServerGroupTestScenario.DEFAULT_TEST_ID,
    'TEST_APP': 'servergrouptest' + ServerGroupTestScenario.DEFAULT_TEST_ID
  }

  return st.ScenarioTestRunner.main(
      ServerGroupTestScenario,
      default_binding_overrides=defaults,
      test_case_list=[ServerGroupTest])

if __name__ == '__main__':
  sys.exit(main())
<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Emanuel Cino <ecino@compassion.ch>
#
#    The licence is in the file __openerp__.py
#
##############################################################################

from openerp.osv import orm
from openerp.tools.translate import _
from openerp.tools import DEFAULT_SERVER_DATE_FORMAT as DF

from sponsorship_compassion.model.product import GIFT_TYPES
from datetime import datetime
from dateutil.relativedelta import relativedelta

from . import gp_connector

class contracts(orm.Model):
    _inherit = 'recurring.contract'

    def create(self, cr, uid, vals, context=None):
        """ When contract is created, push it to GP so that the mailing
        module can access all information. """
        contract_id = super(contracts, self).create(cr, uid, vals, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        # Read data in english
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['lang'] = 'en_US'
        # Write contract in GP
        contract = self.browse(cr, uid, contract_id, ctx)
        self._write_contract_in_gp(uid, gp_connect, contract)
        # Close MySQL connection
        del(gp_connect)

        return contract_id

    def write(self, cr, uid, ids, vals, context=None):
        """ Keep GP updated when a contract is modified. """
        # Read data in english
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['lang'] = 'en_US'
        # Do nothing during the invoice generation process
        if context.get('invoice_generation'):
            return super(contracts, self).write(cr, uid, ids, vals, context)

        ids = [ids] if not isinstance(ids, list) else ids
        gp_connect = gp_connector.GPConnect(cr, uid)
        # If we change the next invoice date, it means we cancel
        # invoices generation and should thus update the situation
        # in GP (advance the months paid).
        if vals.get('next_invoice_date'):
            new_date = datetime.strptime(vals['next_invoice_date'], DF)
            for contract in self.browse(cr, uid, ids, context=ctx):
                old_date = datetime.strptime(contract.next_invoice_date, DF)
                month_diff = relativedelta(new_date, old_date).months
                if contract.state in ('active', 'waiting') and month_diff > 0:
                    if not gp_connect.register_payment(contract.id,
                                                       amount=month_diff):
                        raise orm.except_orm(
                            _("GP Sync Error"),
                            _("Please contact an IT person."))

        # Update GP
        res = super(contracts, self).write(cr, uid, ids, vals, context)
        for contract in self.browse(cr, uid, ids, ctx):
            self._write_contract_in_gp(uid, gp_connect, contract)
        # Close MySQL connection
        del(gp_connect)
        return res

    def _write_contract_in_gp(self, uid, gp_connect, contract):
        if self._is_gp_compatible(contract):
            if not gp_connect.create_or_update_contract(uid, contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("Please contact an IT person."))
        else:
            raise orm.except_orm(
                _("Not compatible with GP"),
                _("You selected some products that are not available "
                  "in GP.") + _("You cannot save this contract."))
        del(gp_connect)

    def _is_gp_compatible(self, contract):
        """ Tells if the contract is compatible with GP. """
        compatible = True
        for line in contract.contract_line_ids:
            compatible = compatible and (
                'Sponsorship' == line.product_id.name
                or line.product_id.gp_fund_id > 0)
        return compatible

    def contract_waiting(self, cr, uid, ids, context=None):
        """ When contract is validated, calculate which month is due
        and push it to GP.
        """
        super(contracts, self).contract_waiting(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.validate_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be validated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def contract_cancelled(self, cr, uid, ids, context=None):
        """ When contract is cancelled, update it in GP. """
        super(contracts, self).contract_cancelled(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.finish_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be terminated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def contract_terminated(self, cr, uid, ids, context=None):
        """ When contract is terminated, update it in GP. """
        super(contracts, self).contract_terminated(cr, uid, ids)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.finish_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be terminated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def _on_contract_active(self, cr, uid, ids, context=None):
        """ When contract is active, update it in GP. """
        super(contracts, self)._on_contract_active(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.activate_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be activated.") +
                    _("Please contact an IT person."))
        del(gp_connect)

    def _invoice_paid(self, cr, uid, invoice, context=None):
        """ When a customer invoice is paid, synchronize GP. """
        super(contracts, self)._invoice_paid(cr, uid, invoice, context)
        if invoice.type == 'out_invoice':
            gp_connect = gp_connector.GPConnect(cr, uid)
            last_pay_date = max([move_line.date
                                 for move_line in invoice.payment_ids
                                 if move_line.credit > 0] or [False])
            contract_ids = set()
            for line in invoice.invoice_line:
                if last_pay_date:
                    gp_connect.insert_affectat(uid, line, last_pay_date)
                else:   # Invoice will go back in open state
                    gp_connect.remove_affectat(invoice.id, line.due_date)
                contract = line.contract_id
                if contract:
                    to_update = (line.product_id.name not in
                                 GIFT_TYPES) and (contract.id
                                                  not in contract_ids)
                    if last_pay_date and to_update:
                        contract_ids.add(contract.id)
                        if not gp_connect.register_payment(contract.id,
                                                           last_pay_date):
                            raise orm.except_orm(
                                _("GP Sync Error"),
                                _("The payment could not be registered into "
                                  "GP.") + _("Please contact an IT person."))
                    elif to_update:
                        contract_ids.add(contract.id)
                        if not gp_connect.undo_payment(contract.id):
                            raise orm.except_orm(
                                _("GP Sync Error"),
                                _("The payment could not be removed from GP.")
                                + _("Please contact an IT person."))
            del(gp_connect)

    def unlink(self, cr, uid, ids, context=None):
        super(contracts, self).unlink(cr, uid, ids, context)
        if not isinstance(ids, list):
            ids = [ids]
        gp_connect = gp_connector.GPConnect(cr, uid)
        gp_connect.delete_contracts(ids)
        del(gp_connect)
        return True

    def reset_open_invoices(self, cr, uid, ids, context=None):
        """Set back GP month status of cancelled invoices because they
        will be replaced by new generated invoices."""
        amount = super(contracts, self).reset_open_invoices(
            cr, uid, ids, context)
        inv_line_obj = self.pool.get('account.invoice.line')
        inv_line_ids = inv_line_obj.search(cr, uid, [
            ('contract_id', 'in', ids),
            ('due_date', '>', datetime.today().strftime(DF)),
            ('state', '=', 'cancel')], context=context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        contract_ids = set()
        for line in inv_line_obj.browse(cr, uid, inv_line_ids, context):
            contract = line.contract_id
            if contract and contract.id not in contract_ids:
                contract_ids.add(contract.id)
                if not gp_connect.undo_payment(contract.id, amount):
                    raise orm.except_orm(
                        _("GP Sync Error"),
                        _("Please contact an IT person."))
        del(gp_connect)

class contract_group(orm.Model):

    """ Update all contracts when group is changed. """
    _inherit = 'recurring.contract.group'

    def write(self, cr, uid, ids, vals, context=None):
        res = super(contract_group, self).write(cr, uid, ids, vals, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for group in self.browse(cr, uid, ids, context):
            for contract in group.contract_ids:
                if not gp_connect.create_or_update_contract(uid, contract):
                    raise orm.except_orm(
                        _("GP Sync Error"),
                        _("Please contact an IT person."))
        return res

    def _generate_invoice_lines(self, cr, uid, contract, invoice_id,
                                context=None):
        """Add in context the information that the generation process is
        working so that GP is not updated during it."""
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['invoice_generation'] = True
        ctx['lang'] = 'en_US'   # Generate everything in english
        super(contract_group, self)._generate_invoice_lines(
            cr, uid, contract, invoice_id, ctx)
<commit_msg>small bug fix<commit_after># -*- encoding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014 Compassion CH (http://www.compassion.ch)
#    Releasing children from poverty in Jesus' name
#    @author: Emanuel Cino <ecino@compassion.ch>
#
#    The licence is in the file __openerp__.py
#
##############################################################################

from openerp.osv import orm
from openerp.tools.translate import _
from openerp.tools import DEFAULT_SERVER_DATE_FORMAT as DF

from sponsorship_compassion.model.product import GIFT_TYPES
from datetime import datetime
from dateutil.relativedelta import relativedelta

from . import gp_connector

class contracts(orm.Model):
    _inherit = 'recurring.contract'

    def create(self, cr, uid, vals, context=None):
        """ When contract is created, push it to GP so that the mailing
        module can access all information. """
        contract_id = super(contracts, self).create(cr, uid, vals, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        # Read data in english
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['lang'] = 'en_US'
        # Write contract in GP
        contract = self.browse(cr, uid, contract_id, ctx)
        self._write_contract_in_gp(uid, gp_connect, contract)
        # Close MySQL connection
        del(gp_connect)

        return contract_id

    def write(self, cr, uid, ids, vals, context=None):
        """ Keep GP updated when a contract is modified. """
        # Read data in english
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['lang'] = 'en_US'
        # Do nothing during the invoice generation process
        if context.get('invoice_generation'):
            return super(contracts, self).write(cr, uid, ids, vals, context)

        ids = [ids] if not isinstance(ids, list) else ids
        gp_connect = gp_connector.GPConnect(cr, uid)
        # If we change the next invoice date, it means we cancel
        # invoices generation and should thus update the situation
        # in GP (advance the months paid).
        if vals.get('next_invoice_date'):
            new_date = datetime.strptime(vals['next_invoice_date'], DF)
            for contract in self.browse(cr, uid, ids, context=ctx):
                old_date = datetime.strptime(contract.next_invoice_date, DF)
                month_diff = relativedelta(new_date, old_date).months
                if contract.state in ('active', 'waiting') and month_diff > 0:
                    if not gp_connect.register_payment(contract.id,
                                                       amount=month_diff):
                        raise orm.except_orm(
                            _("GP Sync Error"),
                            _("Please contact an IT person."))

        # Update GP
        res = super(contracts, self).write(cr, uid, ids, vals, context)
        for contract in self.browse(cr, uid, ids, ctx):
            self._write_contract_in_gp(uid, gp_connect, contract)
        # Close MySQL connection
        del(gp_connect)
        return res

    def _write_contract_in_gp(self, uid, gp_connect, contract):
        if self._is_gp_compatible(contract):
            if not gp_connect.create_or_update_contract(uid, contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("Please contact an IT person."))
        else:
            raise orm.except_orm(
                _("Not compatible with GP"),
                _("You selected some products that are not available "
                  "in GP.") + _("You cannot save this contract."))
        del(gp_connect)

    def _is_gp_compatible(self, contract):
        """ Tells if the contract is compatible with GP. """
        compatible = True
        for line in contract.contract_line_ids:
            compatible = compatible and (
                'Sponsorship' == line.product_id.name
                or 'LDP Sponsorship' == line.product_id.name
                or line.product_id.gp_fund_id > 0)
        return compatible

    def contract_waiting(self, cr, uid, ids, context=None):
        """ When contract is validated, calculate which month is due
        and push it to GP.
        """
        super(contracts, self).contract_waiting(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.validate_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be validated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def contract_cancelled(self, cr, uid, ids, context=None):
        """ When contract is cancelled, update it in GP. """
        super(contracts, self).contract_cancelled(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.finish_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be terminated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def contract_terminated(self, cr, uid, ids, context=None):
        """ When contract is terminated, update it in GP. """
        super(contracts, self).contract_terminated(cr, uid, ids)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.finish_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be terminated.") +
                    _("Please contact an IT person."))
        del(gp_connect)
        return True

    def _on_contract_active(self, cr, uid, ids, context=None):
        """ When contract is active, update it in GP. """
        super(contracts, self)._on_contract_active(cr, uid, ids, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for contract in self.browse(cr, uid, ids, context):
            if not gp_connect.activate_contract(contract):
                raise orm.except_orm(
                    _("GP Sync Error"),
                    _("The sponsorship could not be activated.") +
                    _("Please contact an IT person."))
        del(gp_connect)

    def _invoice_paid(self, cr, uid, invoice, context=None):
        """ When a customer invoice is paid, synchronize GP. """
        super(contracts, self)._invoice_paid(cr, uid, invoice, context)
        if invoice.type == 'out_invoice':
            gp_connect = gp_connector.GPConnect(cr, uid)
            last_pay_date = max([move_line.date
                                 for move_line in invoice.payment_ids
                                 if move_line.credit > 0] or [False])
            contract_ids = set()
            for line in invoice.invoice_line:
                if last_pay_date:
                    gp_connect.insert_affectat(uid, line, last_pay_date)
                else:   # Invoice will go back in open state
                    gp_connect.remove_affectat(invoice.id, line.due_date)
                contract = line.contract_id
                if contract:
                    to_update = (line.product_id.name not in
                                 GIFT_TYPES) and (contract.id
                                                  not in contract_ids)
                    if last_pay_date and to_update:
                        contract_ids.add(contract.id)
                        if not gp_connect.register_payment(contract.id,
                                                           last_pay_date):
                            raise orm.except_orm(
                                _("GP Sync Error"),
                                _("The payment could not be registered into "
                                  "GP.") + _("Please contact an IT person."))
                    elif to_update:
                        contract_ids.add(contract.id)
                        if not gp_connect.undo_payment(contract.id):
                            raise orm.except_orm(
                                _("GP Sync Error"),
                                _("The payment could not be removed from GP.")
                                + _("Please contact an IT person."))
            del(gp_connect)

    def unlink(self, cr, uid, ids, context=None):
        super(contracts, self).unlink(cr, uid, ids, context)
        if not isinstance(ids, list):
            ids = [ids]
        gp_connect = gp_connector.GPConnect(cr, uid)
        gp_connect.delete_contracts(ids)
        del(gp_connect)
        return True

    def reset_open_invoices(self, cr, uid, ids, context=None):
        """Set back GP month status of cancelled invoices because they
        will be replaced by new generated invoices."""
        amount = super(contracts, self).reset_open_invoices(
            cr, uid, ids, context)
        inv_line_obj = self.pool.get('account.invoice.line')
        inv_line_ids = inv_line_obj.search(cr, uid, [
            ('contract_id', 'in', ids),
            ('due_date', '>', datetime.today().strftime(DF)),
            ('state', '=', 'cancel')], context=context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        contract_ids = set()
        for line in inv_line_obj.browse(cr, uid, inv_line_ids, context):
            contract = line.contract_id
            if contract and contract.id not in contract_ids:
                contract_ids.add(contract.id)
                if not gp_connect.undo_payment(contract.id, amount):
                    raise orm.except_orm(
                        _("GP Sync Error"),
                        _("Please contact an IT person."))
        del(gp_connect)

class contract_group(orm.Model):

    """ Update all contracts when group is changed. """
    _inherit = 'recurring.contract.group'

    def write(self, cr, uid, ids, vals, context=None):
        res = super(contract_group, self).write(cr, uid, ids, vals, context)
        gp_connect = gp_connector.GPConnect(cr, uid)
        for group in self.browse(cr, uid, ids, context):
            for contract in group.contract_ids:
                if not gp_connect.create_or_update_contract(uid, contract):
                    raise orm.except_orm(
                        _("GP Sync Error"),
                        _("Please contact an IT person."))
        return res

    def _generate_invoice_lines(self, cr, uid, contract, invoice_id,
                                context=None):
        """Add in context the information that the generation process is
        working so that GP is not updated during it."""
        if context is None:
            context = {}
        ctx = context.copy()
        ctx['invoice_generation'] = True
        ctx['lang'] = 'en_US'   # Generate everything in english
        super(contract_group, self)._generate_invoice_lines(
            cr, uid, contract, invoice_id, ctx)
<|endoftext|>


<commit_before>import uuid

class HandlerMixin(object):
    """
    Mix this in over a ``RequestHandler`` for a correlating header.

    :keyword str correlation_header: the name of the header to use
        for correlation.  If this keyword is omitted, then the header
        is named ``Correlation-ID``.

    This mix-in ensures that responses include a header that correlates
    requests and responses.  If there header is set on the incoming
    request, then it will be copied to the outgoing response.  Otherwise,
    a new UUIDv4 will be generated and inserted.  The value can be
    examined or modified via the ``correlation_id`` property.

    The MRO needs to contain something that resembles a standard
    :class:`tornado.web.RequestHandler`.  Specifically, we need the
    following things to be available:

    - :meth:`~tornado.web.RequestHandler.prepare` needs to be called
      appropriately
    - :meth:`~tornado.web.RequestHandler.set_header` needs to exist in
      the MRO and it needs to overwrite the header value
    - :meth:`~tornado.web.RequestHandler.set_default_headers` should be
      called to establish the default header values
    - ``self.request`` is a object that has a ``headers`` property that
      contains the request headers as a ``dict``.

    """

    def __init__(self, *args, **kwargs):
        # correlation_id is used from within set_default_headers
        # which is called from within super().__init__() so we need
        # to make sure that it is set *BEFORE* we call super.
        self.__header_name = kwargs.pop(
            'correlation_header', 'Correlation-ID')
        self.__correlation_id = str(uuid.uuid4())
        super(HandlerMixin, self).__init__(*args, **kwargs)

    def prepare(self):
        # Here we want to copy an incoming Correlation-ID header if
        # one exists.  We also want to set it in the outgoing response
        # which the property setter does for us.
        super(HandlerMixin, self).prepare()
        correlation_id = self.get_request_header(self.__header_name, None)
        if correlation_id is not None:
            self.correlation_id = correlation_id

    def set_default_headers(self):
        # This is called during initialization as well as *AFTER*
        # calling clear() when an error occurs so we need to make
        # sure that our header is set again...
        super(HandlerMixin, self).set_default_headers()
        self.set_header(self.__header_name, self.correlation_id)

    @property
    def correlation_id(self):
        """Correlation header value."""
        return self.__correlation_id

    @correlation_id.setter
    def correlation_id(self, value):
        self.__correlation_id = value
        self.set_header(self.__header_name, self.__correlation_id)

    def get_request_header(self, name, default):
        """
        Retrieve the value of a request header.

        :param str name: the name of the header to retrieve
        :param default: the value to return if the header is not set

        This method abstracts the act of retrieving a header value out
        from the implementation.  This makes it possible to implement
        a *RequestHandler* that is something other than a
        :class:`tornado.web.RequestHandler` by simply implementing this
        method and ``set_header`` over the underlying implementation,
        for example, say AMQP message properties.

        """
        return self.request.headers.get(name, default)
<commit_msg>Create new log_function to add CID to tornado logs<commit_after>import uuid

class HandlerMixin(object):
    """
    Mix this in over a ``RequestHandler`` for a correlating header.

    :keyword str correlation_header: the name of the header to use
        for correlation.  If this keyword is omitted, then the header
        is named ``Correlation-ID``.

    This mix-in ensures that responses include a header that correlates
    requests and responses.  If there header is set on the incoming
    request, then it will be copied to the outgoing response.  Otherwise,
    a new UUIDv4 will be generated and inserted.  The value can be
    examined or modified via the ``correlation_id`` property.

    The MRO needs to contain something that resembles a standard
    :class:`tornado.web.RequestHandler`.  Specifically, we need the
    following things to be available:

    - :meth:`~tornado.web.RequestHandler.prepare` needs to be called
      appropriately
    - :meth:`~tornado.web.RequestHandler.set_header` needs to exist in
      the MRO and it needs to overwrite the header value
    - :meth:`~tornado.web.RequestHandler.set_default_headers` should be
      called to establish the default header values
    - ``self.request`` is a object that has a ``headers`` property that
      contains the request headers as a ``dict``.

    """

    def __init__(self, *args, **kwargs):
        # correlation_id is used from within set_default_headers
        # which is called from within super().__init__() so we need
        # to make sure that it is set *BEFORE* we call super.
        self.__header_name = kwargs.pop(
            'correlation_header', 'Correlation-ID')
        self.__correlation_id = str(uuid.uuid4())
        super(HandlerMixin, self).__init__(*args, **kwargs)

    def prepare(self):
        # Here we want to copy an incoming Correlation-ID header if
        # one exists.  We also want to set it in the outgoing response
        # which the property setter does for us.
        super(HandlerMixin, self).prepare()
        correlation_id = self.get_request_header(self.__header_name, None)
        if correlation_id is not None:
            self.correlation_id = correlation_id

    def set_default_headers(self):
        # This is called during initialization as well as *AFTER*
        # calling clear() when an error occurs so we need to make
        # sure that our header is set again...
        super(HandlerMixin, self).set_default_headers()
        self.set_header(self.__header_name, self.correlation_id)

    @property
    def correlation_id(self):
        """Correlation header value."""
        return self.__correlation_id

    @correlation_id.setter
    def correlation_id(self, value):
        self.__correlation_id = value
        self.set_header(self.__header_name, self.__correlation_id)

    def get_request_header(self, name, default):
        """
        Retrieve the value of a request header.

        :param str name: the name of the header to retrieve
        :param default: the value to return if the header is not set

        This method abstracts the act of retrieving a header value out
        from the implementation.  This makes it possible to implement
        a *RequestHandler* that is something other than a
        :class:`tornado.web.RequestHandler` by simply implementing this
        method and ``set_header`` over the underlying implementation,
        for example, say AMQP message properties.

        """
        return self.request.headers.get(name, default)

def correlation_id_logger(handler):
    """ Custom Tornado access log writer that appends correlation-id.

    This function can be used to append the coorelation-id to the
    Tornado access logs. To use, simply set the value of the
    log_function kwarg of the Tornado Application constructor to this
    function.

    *Example*
    web.Application([], log_function=correlation_id_logger)

    :param tornado.web.RequestHandler handler: the request handler that
        is processing the client request.
    """
    if handler.get_status() < 400:
        log_method = log.access_log.info
    elif handler.get_status() < 500:
        log_method = log.access_log.warning
    else:
        log_method = log.access_log.error
    request_time = 1000.0 * handler.request.request_time()
    correlation_id = getattr(handler, "correlation_id", None)
    if correlation_id is None:
        correlation_id = handler.request.headers.get('Correlation-ID', None)
    log_method("%d %s %.2fms {CID %s}", handler.get_status(),
               handler._request_summary(), request_time, correlation_id)
<|endoftext|>


<commit_before>from pages.models import HomePage, HigherEducation, GeneralPage, ContactUs, AboutUs

from wagtail.contrib.wagtailapi.serializers import PageSerializer

class HomePageSerializer(PageSerializer):
    class Meta:
        model = HomePage
        fields = ('id',
                  'title',
                  'row_1',
                  'row_2',
                  'row_3',
                  'row_4',
                  'row_5',
                  'slug',
                  'seo_title',
                  'search_description',
                  )

class HigherEducationSerializer(PageSerializer):
    class Meta:
        model = HigherEducation
        fields = ('id',
                  'title',
                  'intro_heading',
                  'intro_description',
                  'row_1',
                  'get_started_heading',
                  'get_started_step_1_heading',
                  'get_started_step_1_description',
                  'get_started_step_1_cta',
                  'get_started_step_2_heading',
                  'get_started_step_2_description',
                  'get_started_step_2_cta',
                  'get_started_step_3_heading',
                  'get_started_step_3_description',
                  'get_started_step_3_cta',
                  'adopt_heading',
                  'adopt_description',
                  'adopt_cta',
                  'row_2',
                  'row_3',
                  'slug',
                  'seo_title',
                  'search_description',
                  )

class GeneralPageSerializer(PageSerializer):
    class Meta:
        model = GeneralPage
        fields = ('id',
                  'title',
                  'body',
                  'slug',
                  'seo_title',
                  'search_description',
                  )

class ContactUsSerializer(PageSerializer):
    class Meta:
        model = ContactUs
        fields = ('id',
                  'title',
                  'tagline',
                  'mailing_header',
                  'mailing_address',
                  'phone_number',
                  'slug',
                  'seo_title',
                  'search_description',
                  )

class AboutUsSerializer(PageSerializer):
    class Meta:
        model = AboutUs
        fields = ('tagline',
                  'intro_heading',
                  'intro_paragraph',
                  'our_team_heading',
                  'openstax_team',
                  'strategic_advisors',
                  'slug',
                  'seo_title',
                  'search_description',
                  )<commit_msg>use model defined api_fields on serializers<commit_after>from pages.models import HomePage, HigherEducation, GeneralPage, ContactUs, AboutUs

from wagtail.contrib.wagtailapi.serializers import PageSerializer

class HomePageSerializer(PageSerializer):
    class Meta:
        model = HomePage
        fields = HomePage.api_fields

class HigherEducationSerializer(PageSerializer):
    class Meta:
        model = HigherEducation
        fields = HigherEducation.api_fields

class GeneralPageSerializer(PageSerializer):
    class Meta:
        model = GeneralPage
        fields = GeneralPage.api_fields

class ContactUsSerializer(PageSerializer):
    class Meta:
        model = ContactUs
        fields = ContactUs.api_fields

class AboutUsSerializer(PageSerializer):
    class Meta:
        model = AboutUs
        fields = AboutUs.api_fields
<|endoftext|>