From f3389067360b70aea196c305b23048ca3b2aed0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=92=D0=B8=D0=BA=D1=82=D0=BE=D1=80=20=D0=93=D0=BB=D0=B0?=
 =?UTF-8?q?=D0=B4=D0=BA=D0=B8=D1=85?= <v.gladkikh@spider.ru>
Date: Mon, 21 Oct 2019 11:44:47 +0300
Subject: [PATCH] Fix location tests

---
 apps/location/models.py     |  4 ++++
 apps/location/tests.py      | 25 +++++++++++++++++++++++++
 apps/location/views/back.py |  4 +++-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/apps/location/models.py b/apps/location/models.py
index 7f797811..2b7aa363 100644
--- a/apps/location/models.py
+++ b/apps/location/models.py
@@ -21,6 +21,10 @@ class Country(TranslatedFieldsMixin, SVGImageMixin, ProjectBaseMixin):
     high_price = models.IntegerField(default=50, verbose_name=_('High price'))
     languages = models.ManyToManyField(Language, verbose_name=_('Languages'))
 
+    @property
+    def country_id(self):
+        return self.id
+
     class Meta:
         """Meta class."""
 
diff --git a/apps/location/tests.py b/apps/location/tests.py
index 0770227a..edb719bd 100644
--- a/apps/location/tests.py
+++ b/apps/location/tests.py
@@ -59,15 +59,25 @@ class CountryTests(BaseTestCase):
     def setUp(self):
         super().setUp()
 
+
     def test_country_CRUD(self):
         data = {
             'name': {"ru-RU":"Russia"},
             'code': 'test'
         }
+
         response = self.client.post('/api/back/location/countries/', data=data, format='json')
         response_data = response.json()
         self.assertEqual(response.status_code, status.HTTP_201_CREATED)
 
+        country = Country.objects.get(pk=response_data["id"])
+        role = Role.objects.create(role=Role.COUNTRY_ADMIN, country=country)
+        role.save()
+
+        user_role = UserRole.objects.create(user=self.user, role=role)
+
+        user_role.save()
+
         response = self.client.get('/api/back/location/countries/', format='json')
         self.assertEqual(response.status_code, status.HTTP_200_OK)
 
@@ -94,6 +104,14 @@ class RegionTests(BaseTestCase):
             code="test"
         )
 
+        role = Role.objects.create(role=Role.COUNTRY_ADMIN, country=self.country)
+        role.save()
+
+        user_role = UserRole.objects.create(user=self.user, role=role)
+
+        user_role.save()
+
+
     def test_region_CRUD(self):
         response = self.client.get('/api/back/location/regions/', format='json')
         self.assertEqual(response.status_code, status.HTTP_200_OK)
@@ -138,6 +156,13 @@ class CityTests(BaseTestCase):
             country=self.country
         )
 
+        role = Role.objects.create(role=Role.COUNTRY_ADMIN, country=self.country)
+        role.save()
+
+        user_role = UserRole.objects.create(user=self.user, role=role)
+
+        user_role.save()
+
     def test_city_CRUD(self):
         response = self.client.get('/api/back/location/cities/', format='json')
         self.assertEqual(response.status_code, status.HTTP_200_OK)
diff --git a/apps/location/views/back.py b/apps/location/views/back.py
index 3a2739b2..cb8246a4 100644
--- a/apps/location/views/back.py
+++ b/apps/location/views/back.py
@@ -47,6 +47,7 @@ class RegionRUDView(common.RegionViewMixin, generics.RetrieveUpdateDestroyAPIVie
 # Country
 class CountryListCreateView(generics.ListCreateAPIView):
     """List/Create view for model Country."""
+    queryset = models.Country.objects.all()
     serializer_class = serializers.CountryBackSerializer
     pagination_class = None
     permission_classes = [IsCountryAdmin]
@@ -54,4 +55,5 @@ class CountryListCreateView(generics.ListCreateAPIView):
 class CountryRUDView(generics.RetrieveUpdateDestroyAPIView):
     """RUD view for model Country."""
     serializer_class = serializers.CountryBackSerializer
-    permission_classes = [IsCountryAdmin]
\ No newline at end of file
+    permission_classes = [IsCountryAdmin]
+    queryset = models.Country.objects.all()
\ No newline at end of file