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