Fix find-basket for own request

1 files changed, 23 insertions, 4 deletions

diff --git a/ishtar_common/tests.py b/ishtar_common/tests.py
index 9e5fc3466..c128dd93a 100644
--- a/ishtar_common/tests.py
+++ b/ishtar_common/tests.py
@@ -626,15 +626,18 @@ class ShortMenuTest(TestCase):
         self.assertEqual(response.status_code, 200)
         self.assertFalse(str(tf.cached_label) in response.content)
 
-    def testTreatment(self):
-        c = Client()
-        c.login(username=self.username, password=self.password)
+    def _create_treatment(self):
         from archaeological_finds.models import Treatment
-        treat = Treatment.objects.create(
+        return Treatment.objects.create(
             label="My treatment",
             year=2052
         )
 
+    def testTreatment(self):
+        c = Client()
+        c.login(username=self.username, password=self.password)
+        treat = self._create_treatment()
+
         # not available at first
         response = c.get(reverse('shortcut-menu'))
         self.assertEqual(response.status_code, 200)
@@ -684,6 +687,22 @@ class ShortMenuTest(TestCase):
         self.assertFalse(c.session['find'])
         self.assertFalse(c.session['contextrecord'])
 
+        self._create_treatment()
+
+    def test_basket_hierarchy_with_treatments(self):
+        c = Client()
+        c.login(username=self.username, password=self.password)
+        from archaeological_finds.models import FindBasket
+        basket = FindBasket.objects.create(
+            label="My basket",
+            user=self.user.ishtaruser
+        )
+        session = c.session
+        session['find'] = 'basket-{}'.format(basket.pk)
+        session.save()
+        response = c.get(reverse('shortcut-menu'))
+        self.assertEqual(response.status_code, 200)
+
 
 class ImportTest(TestCase):
     def testDeleteRelated(self):