Django Rest Framework Test: How to refresh a relation? - django

Question

I have a permission in Django Rest Framework:
from annoying.functions import get_object_or_None
from django.utils.translation import ugettext_lazy as _
from rest_framework import permissions
from restaurants.models import Restaurant
class TableBelongsToRestaurantPermission(permissions.BasePermission):
"""
Permission to check if the table belongs to the restaurant in the request. This ensures
(together with the UserOwnsRestaurantPermission) that owner can change the QR code of a
restaurant that he doesn't own.
"""
message = TABLE_BELONGS_TO_RESTAURANT_PERMISSION_DENIED_MESSAGE
def has_object_permission(self, request, view, obj):
if not obj.table:
return True
slug = request.data.get("restaurant_slug", "")
restaurant = get_object_or_None(Restaurant, slug=slug)
if restaurant:
return restaurant.table_set.filter(id=obj.table.id).exists()
return False
And now, I wrote tests for this:
from unittest import mock
from allauth.account.models import EmailAddress
from django.contrib.auth import get_user_model
from django.test import TestCase
from addresses.models import (Address, City, Country, PostalCode, State,
StreetName)
from core.utils import QR_CODE_FUNCTIONS
from employments.models import Employment
from licenses.models import RestaurantLicense
from profiles.models import UserOwnerProfile
from qrcodeproperties.models import QRCodePropertyCheckinUser
from restaurants.models import Restaurant
from tables.models import Table
from ...models import QRCode, QRCodeFunction
from ..permissions import TableBelongsToRestaurantPermission
User = get_user_model()
USER_OWNER_EMAIL = "owner#burgergrill.de"
USER_OWNER_NAME = "Owner"
USER_PASSWORD = "test1234test"
ISO_ALPHA_2_CODE = "DE"
STATE_NAME = "NRW"
CITY_NAME = "Köln"
STREET_NAME = "Burgerstraße"
POSTAL_CODE = "32062"
STREET_NUMBER = "119"
RESTAURANT_NAME = "Burgergrill"
RESTAURANT_SLUG = "burgergrill"
TABLE_NUMBER = 1
OTHER_RESTAURANT_NAME = "Extrablatt"
OTHER_RESTAURANT_SLUG = "extrablatt"
class TableBelongsToRestaurantPermissionTestCase(TestCase):
def setUp(self):
self.permission = TableBelongsToRestaurantPermission()
owner_user = User.objects.create_user(
email=USER_OWNER_EMAIL,
name=USER_OWNER_NAME,
password=USER_PASSWORD
)
self.owner_user = owner_user
owner = UserOwnerProfile.objects.create(user=owner_user)
owner.save()
self.owner = owner
emailaddress = EmailAddress.objects.create(
user=owner_user,
email=owner_user.email,
verified=True,
primary=True
)
emailaddress.save()
self.emailaddress = emailaddress
country = Country.objects.create(
country=ISO_ALPHA_2_CODE
)
country.save()
self.country = country
state = State.objects.create(
name=STATE_NAME,
country=self.country
)
state.save()
self.state = state
city = City.objects.create(
name=CITY_NAME,
country=self.country,
state=self.state
)
city.save()
self.city = city
street_name = StreetName.objects.create(
name=STREET_NAME,
city=self.city
)
street_name.save()
self.street_name = street_name
postal_code = PostalCode.objects.create(
code=POSTAL_CODE,
city=self.city
)
postal_code.save()
self.postal_code = postal_code
address = Address.objects.create(
country=self.country,
state=self.state,
city=self.city,
postal_code=self.postal_code,
street_name=self.street_name,
street_number=STREET_NUMBER
)
address.save()
self.address = address
restaurant = Restaurant.objects.create(
name=RESTAURANT_NAME,
address=self.address
)
restaurant.save()
self.restaurant = restaurant
qr_code = QRCode.objects.create()
qr_code.save()
self.qr_code = qr_code
self.request = mock.MagicMock(
user=owner_user,
kwargs={"slug": RESTAURANT_SLUG}
)
self.view = mock.MagicMock()
def test_permissions_table_belongs_to_restaurant_denied_when_table_belongs_to_different_restaurant(self):
"""Test if a user gets permission denied if the qr code's table belongs to another restaurant."""
restaurant = Restaurant.objects.create(name=OTHER_RESTAURANT_NAME, address=self.address)
restaurant.save()
table = Table.objects.create(table_number=TABLE_NUMBER, restaurant=restaurant)
table.save()
prev_count = Table.objects.all().count()
qr_code_function = QRCodeFunction.objects.create(
qr_code=self.qr_code,
qr_code_function=QR_CODE_FUNCTIONS.checkin_user
)
qr_code_function.save()
qr_code_property = QRCodePropertyCheckinUser(
qr_code=self.qr_code,
table=table,
)
qr_code_property.save()
self.assertFalse(self.permission.has_object_permission(
request=self.request, view=self.view, obj=self.qr_code
))
def test_permissions_table_belongs_to_restaurant_granted_when_table_belongs_to_same_restaurant(self):
"""Test if a user gets permission granted if the qr code's table belongs to the same restaurant."""
table = Table.objects.create(table_number=TABLE_NUMBER, restaurant=self.restaurant)
table.save()
prev_count = Table.objects.all().count()
qr_code_function = QRCodeFunction.objects.create(
qr_code=self.qr_code,
qr_code_function=QR_CODE_FUNCTIONS.checkin_user
)
qr_code_function.save()
qr_code_property = QRCodePropertyCheckinUser(
qr_code=self.qr_code,
table=table,
)
qr_code_property.save()
self.restaurant.refresh_from_db()
print(table.restaurant == self.restaurant)
print(self.qr_code.table)
print(self.restaurant.table_set)
self.assertTrue(self.permission.has_object_permission(
request=self.request, view=self.view, obj=self.qr_code
))
def test_permissions_table_belongs_to_restaurant_granted_when_qr_code_without_table(self):
"""Test if a user gets permission granted if the qr code has no table."""
self.assertTrue(self.permission.has_object_permission(
request=self.request, view=self.view, obj=self.qr_code
))
The problem is that the second test (test_permissions_table_belongs_to_restaurant_granted_when_table_belongs_to_same_restaurant) fails. The print statements give out True, 1 Pontgrill and tables.Table.None respectively. That means that for some reason even though I did .refresh_from_db() the data for the restaurant does not update and the test fails because it Django thinks there are no relations from Table to Restaurant.
How can you refresh relations and fix the test?