Cleaning up GFK test module. Refs #607.

2 files changed, 31 insertions, 33 deletions

diff --git a/rest_framework/tests/genericrelations.py b/rest_framework/tests/genericrelations.py
index bc7378e1..dcdd8329 100644
--- a/rest_framework/tests/genericrelations.py
+++ b/rest_framework/tests/genericrelations.py
@@ -1,23 +1,42 @@
+from django.contrib.contenttypes.models import ContentType
+from django.contrib.contenttypes.generic import GenericRelation, GenericForeignKey
+from django.db import models
 from django.test import TestCase
 from rest_framework import serializers
-from rest_framework.tests.models import *
+
+
+class Tag(models.Model):
+    """
+    Tags have a descriptive slug, and are attached to an arbitrary object.
+    """
+    tag = models.SlugField()
+    content_type = models.ForeignKey(ContentType)
+    object_id = models.PositiveIntegerField()
+    content_object = GenericForeignKey('content_type', 'object_id')
+
+    def __unicode__(self):
+        return self.tag
+
+
+class Bookmark(models.Model):
+    """
+    A URL bookmark that may have multiple tags attached.
+    """
+    url = models.URLField()
+    tags = GenericRelation(Tag)
 
 
 class TestGenericRelations(TestCase):
     def setUp(self):
-        bookmark = Bookmark(url='https://www.djangoproject.com/')
-        bookmark.save()
-        django = Tag(tag_name='django')
-        django.save()
-        python = Tag(tag_name='python')
-        python.save()
-        t1 = TaggedItem(content_object=bookmark, tag=django)
-        t1.save()
-        t2 = TaggedItem(content_object=bookmark, tag=python)
-        t2.save()
-        self.bookmark = bookmark
+        self.bookmark = Bookmark.objects.create(url='https://www.djangoproject.com/')
+        Tag.objects.create(content_object=self.bookmark, tag='django')
+        Tag.objects.create(content_object=self.bookmark, tag='python')
 
     def test_reverse_generic_relation(self):
+        """
+        Test a relationship that spans a GenericRelation field.
+        """
+
         class BookmarkSerializer(serializers.ModelSerializer):
             tags = serializers.ManyRelatedField(source='tags')
 
diff --git a/rest_framework/tests/models.py b/rest_framework/tests/models.py
index 93f09761..9ab15328 100644
--- a/rest_framework/tests/models.py
+++ b/rest_framework/tests/models.py
@@ -86,27 +86,6 @@ class ReadOnlyManyToManyModel(RESTFrameworkModel):
     text = models.CharField(max_length=100, default='anchor')
     rel = models.ManyToManyField(Anchor)
 
-# Models to test generic relations
-
-
-class Tag(RESTFrameworkModel):
-    tag_name = models.SlugField()
-
-
-class TaggedItem(RESTFrameworkModel):
-    tag = models.ForeignKey(Tag, related_name='items')
-    content_type = models.ForeignKey(ContentType)
-    object_id = models.PositiveIntegerField()
-    content_object = GenericForeignKey('content_type', 'object_id')
-
-    def __unicode__(self):
-        return self.tag.tag_name
-
-
-class Bookmark(RESTFrameworkModel):
-    url = models.URLField()
-    tags = GenericRelation(TaggedItem)
-
 
 # Model to test filtering.
 class FilterableItem(RESTFrameworkModel):