def test_get_user_org(self):
"""Test getting org from user"""
fake_user = User.objects.create(username='******')
fake_org_1 = Organization.objects.create()
fake_org_2 = Organization.objects.create()
fake_org_3 = Organization.objects.create()
OrganizationUser.objects.create(user=fake_user,
organization=fake_org_1)
OrganizationUser.objects.create(user=fake_user,
organization=fake_org_2)
OrganizationUser.objects.create(user=fake_user,
organization=fake_org_3)
# no default_organization and no parent org
result = get_user_org(fake_user)
self.assertIn(result, fake_user.orgs.all())
# parent org, no default_organization
fake_org_1.parent_org = fake_org_2
fake_org_1.save()
expected = fake_org_2
result = get_user_org(fake_user)
self.assertEqual(result, expected)
# user default_organization
fake_user.default_organization = fake_org_3
fake_user.save()
expected = fake_org_3
result = get_user_org(fake_user)
self.assertEqual(result, expected)
def get_organization(self, request, return_obj=False):
"""Get org from query param or request.user.
:param request: request object.
:param return_obj: bool. Set to True if obj vs pk is desired.
:return: int representing a valid organization pk or
organization object.
"""
# print("my return obj is set to %s" % return_obj)
if not request.user:
return None
if not getattr(self, '_organization', None):
org_id = get_org_id(request)
org = None
if not org_id:
org = get_user_org(request.user)
org_id = int(getattr(org, 'pk'))
if return_obj:
if not org:
try:
org = request.user.orgs.get(pk=org_id)
self._organization = org
except ObjectDoesNotExist:
raise PermissionDenied('Incorrect org id.')
else:
self._organization = org
else:
self._organization = org_id
return self._organization
