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
def test_get_org_id(self):
"""Test getting org id from request."""
mock_request = mock.MagicMock()
mock_request.query_params = {'organization': 1}
result = get_org_id(mock_request)
self.assertEqual(1, result)
mock_request = mock.MagicMock()
mock_request.query_params = {'organization': None}
mock_request.data = {'organization': 2}
result = get_org_id(mock_request)
self.assertEqual(2, result)
mock_request = mock.MagicMock()
mock_request.query_params = {'organization': None}
mock_value_error = mock.PropertyMock(side_effect=ValueError)
type(mock_request).data = mock_value_error
result = get_org_id(mock_request)
self.assertEqual(None, result)
