tests/data_seeder.py

   1 #/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3 #
   4 # Copyright (c) 2010-2012 Cidadania S. Coop. Galega
   5 #
   6 # This file is part of e-cidadania.
   7 #
   8 # e-cidadania is free software: you can redistribute it and/or modify
   9 # it under the terms of the GNU General Public License as published by
  10 # the Free Software Foundation, either version 3 of the License, or
  11 # (at your option) any later version.
  12 #
  13 # e-cidadania is distributed in the hope that it will be useful,
  14 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  15 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  16 # GNU General Public License for more details.
  17 #
  18 # You should have received a copy of the GNU General Public License
  19 # along with e-cidadania. If not, see <http://www.gnu.org/licenses/>.
  20
  21 from mockups import Mockup
  22 from mockups.helpers import register
  23
  24 from django.contrib.auth.models import User
  25 from django.db import models
  26
  27
  28 class Error(Exception):
  29     """Base class for all exceptions raised by the seeder.
  30     """
  31     pass
  32
  33
  34 class InvalidModelError(Error):
  35     """Raised when a model is not a valid Django model.
  36     """
  37     pass
  38
  39
  40 class DataSeeder(object):
  41     """Class which contains the methods to seed data or fixtures for testing.
  42
  43     We presently wrap our methods around Mockup.
  44     """
  45
  46     def __init__(self):
  47         #Register User with Mockup
  48         register(User, Mockup)
  49
  50     def seed(self, model, constraints=None, follow_fk=None, generate_fk=None,
  51              follow_m2m=None, factory=None, model_properties=None, commit=True):
  52         """Creates and saves an instance of 'model' in the database.
  53
  54         The values generated by Mockup class for the fields may not be
  55         acceptable. Custom values for the fields can be provided in
  56         model_properties.
  57
  58         Parameters:
  59             ``model``: A model class which is used to create the test data.
  60
  61             ``constraints``: A list of callables. The constraints are used to
  62             verify if the created model instance may be used. The callable
  63             gets the actual model as first and the instance as second
  64             parameter. The instance is not populated yet at this moment.  The
  65             callable may raise an :exc:`InvalidConstraint` exception to
  66             indicate which fields violate the constraint.
  67
  68             ``follow_fk``: A boolean value indicating if foreign keys should be
  69             set to random, already existing, instances of the related model.
  70
  71             ``generate_fk``: A boolean which indicates if related models should
  72             also be created with random values. The *follow_fk* parameter will
  73             be ignored if *generate_fk* is set to ``True``.
  74
  75             ``follow_m2m``: A tuple containing minium and maximum of model
  76             instances that are assigned to ``ManyToManyField``. No new
  77             instances will be created. Default is (1, 5).  You can ignore
  78             ``ManyToManyField`` fields by setting this parameter to ``False``.
  79
  80             ``generate_m2m``: A tuple containing minimum and maximum number of
  81             model instance that are newly created and assigned to the
  82             ``ManyToManyField``. Default is ``False`` which disables the
  83             generation of new related instances. The value of ``follow_m2m``
  84             will be ignored if this parameter is set.
  85
  86             ``factory``: A Factory *instance*, overriding the one defined in the
  87             Mockup class.
  88
  89             ``model_properties``: A dict containing the custom properties
  90             for the ``model``
  91
  92             ``commit``: A boolean which is set to True by default and indicates
  93             whether the model should be saved to the database or not.
  94
  95         """
  96
  97         #if not isinstance(model, models.Model):
  98         #    raise InvalidModelError("%s is not a valid Django model." % model)
  99         if model_properties is None:
 100             model_properties = {}
 101
 102         # Creates and randomly populates the data
 103         mockup = Mockup(model, constraints=constraints, follow_fk=follow_fk,
 104                         generate_fk=generate_fk, follow_m2m=follow_m2m,
 105                         factory=factory)
 106         created_model = mockup.create_one(commit=commit)
 107
 108         # set the attributes of the model as provided in model_properties
 109         for key in model_properties.iterkeys():
 110             setattr(created_model, key, model_properties[key])
 111         if commit:
 112             created_model.save()
 113         return created_model
 114
 115     def seedn(self, count, model, constraints=None, follow_fk=None,
 116               generate_fk=None, follow_m2m=None, factory=None,
 117               model_properties=None, commit=True ):
 118         """Creates and saves n instances of 'model' in the database and returns
 119         a list of all those saved instances.
 120
 121         The method uses self.seed to generate a list of instances of ``model``
 122         ``count`` number of times.
 123         """
 124
 125         obj_list = []
 126         for _ in xrange(count):
 127             obj = self.seed(model=model, constraints=constraints,
 128                             follow_fk=follow_fk, generate_fk=generate_fk,
 129                             follow_m2m=follow_m2m, factory=factory,
 130                             model_properties=model_properties, commit=commit)
 131             obj_list.append(obj)
 132
 133         return obj_list
 134
 135 seeder = DataSeeder()