Add impersonation APIs, add json kwarg to make_request (#46)

- ServerContext.make_request: payload is now optional
  - Updated usages that passed None to pass nothing
- ServerContext.make_request: add json kwarg
  - This automatically does json_dumps and sets the content-type header for you
  - Updated usages of make_request to use the json kwarg where applicable
- Add impersonate_user and stop_impersonating to security module (including APIWrapper)
- Fix imports for mock in unit tests
- Remove unused test_utils file.
- Add more environment variables for integration test configuration
  - host, port, context_path can all be overridden via env vars

Author: labkey-alan

.gitignore

@@ -13,3 +13,4 @@ labkey\.egg-info/
 
 .DS_Store
 .python-version
+.vscode/

CHANGE.txt

@@ -6,7 +6,13 @@ What's New in the LabKey 2.1.0 package
 ==============================
 
 *Release date: TBD*
-Adding support for ontology based column filters ONTOLOGY_IN_SUBTREE and ONTOLOGY_NOT_IN_SUBTREE
+- Add support for ontology based column filters ONTOLOGY_IN_SUBTREE and ONTOLOGY_NOT_IN_SUBTREE
+- ServerContext.make_request: payload is now optional
+- ServerContext.make_request: add json kwarg
+  - This automatically does json_dumps and sets the content-type header for you
+- Add impersonate_user and stop_impersonating to security module (including APIWrapper)
+- Add more environment variables for integration test configuration
+  - host, port, and context_path can now be overridden via env vars
 
 What's New in the LabKey 2.0.1 package
 ==============================

labkey/container.py

@@ -1,5 +1,4 @@
 from .server_context import ServerContext
-from labkey.utils import json_dumps
 
 
 def create(
@@ -23,7 +22,6 @@ def create(
     :param title: the title for the container.
     :return:
     """
-    headers = {"Content-Type": "application/json"}
     url = server_context.build_url("core", "createContainer.api", container_path)
     payload = {
         "description": description,
@@ -32,7 +30,7 @@ def create(
         "name": name,
         "title": title,
     }
-    return server_context.make_request(url, json_dumps(payload), headers=headers)
+    return server_context.make_request(url, json=payload)
 
 
 def delete(server_context: ServerContext, container_path: str = None) -> any:
@@ -44,7 +42,7 @@ def delete(server_context: ServerContext, container_path: str = None) -> any:
     """
     headers = {"Content-Type": "application/json"}
     url = server_context.build_url("core", "deleteContainer.api", container_path)
-    return server_context.make_request(url, None, headers=headers)
+    return server_context.make_request(url, headers=headers)
 
 
 class ContainerWrapper:

labkey/domain.py

@@ -17,7 +17,6 @@
 from typing import Union, List
 
 from .server_context import ServerContext
-from labkey.utils import json_dumps
 from labkey.query import QueryFilter
 
 
@@ -390,13 +389,12 @@ def create(
     :return: Domain
     """
     url = server_context.build_url("property", "createDomain.api", container_path=container_path)
-    headers = {"Content-Type": "application/json"}
     domain = None
     domain_fields = domain_definition["domainDesign"]["fields"]
     domain_definition["domainDesign"]["fields"] = list(
         map(__format_conditional_filters, domain_fields)
     )
-    raw_domain = server_context.make_request(url, json_dumps(domain_definition), headers=headers)
+    raw_domain = server_context.make_request(url, json=domain_definition)
 
     if raw_domain is not None:
         domain = Domain(**raw_domain)
@@ -416,10 +414,9 @@ def drop(
     :return:
     """
     url = server_context.build_url("property", "deleteDomain.api", container_path=container_path)
-    headers = {"Content-Type": "application/json"}
     payload = {"schemaName": schema_name, "queryName": query_name}
 
-    return server_context.make_request(url, json_dumps(payload), headers=headers)
+    return server_context.make_request(url, json=payload)
 
 
 def get(
@@ -455,7 +452,7 @@ def infer_fields(
     :return:
     """
     url = server_context.build_url("property", "inferDomain.api", container_path=container_path)
-    raw_infer = server_context.make_request(url, None, file_payload={"inferfile": data_file})
+    raw_infer = server_context.make_request(url, file_payload={"inferfile": data_file})
 
     fields = None
     if "fields" in raw_infer:
@@ -483,14 +480,13 @@ def save(
     :return:
     """
     url = server_context.build_url("property", "saveDomain.api", container_path=container_path)
-    headers = {"Content-Type": "application/json"}
     payload = {
         "domainDesign": domain.to_json(),
         "queryName": query_name,
         "schemaName": schema_name,
     }
 
-    return server_context.make_request(url, json_dumps(payload), headers=headers)
+    return server_context.make_request(url, json=payload)
 
 
 class DomainWrapper:

labkey/experiment.py

@@ -17,7 +17,6 @@
 from typing import List, Optional
 
 from .server_context import ServerContext
-from labkey.utils import json_dumps
 
 
 class ExpObject:
@@ -103,7 +102,7 @@ def to_json(self):
         data["materialOutputs"] = self.material_outputs
         data["plateMetadata"] = self.plate_metadata
 
-        # Issue 2489: Drop empty values. Server supplies default values for missing keys, 
+        # Issue 2489: Drop empty values. Server supplies default values for missing keys,
         # and will throw exception if a null value is supplied
         data = {k: v for k, v in data.items() if v}
         return data
@@ -163,14 +162,9 @@ def load_batch(server_context: ServerContext, assay_id: int, batch_id: int) -> O
     """
     load_batch_url = server_context.build_url("assay", "getAssayBatch.api")
     loaded_batch = None
-
     payload = {"assayId": assay_id, "batchId": batch_id}
+    json_body = server_context.make_request(load_batch_url, json=payload)
 
-    headers = {"Content-type": "application/json", "Accept": "text/plain"}
-
-    json_body = server_context.make_request(
-        load_batch_url, json_dumps(payload, sort_keys=True), headers=headers
-    )
     if json_body is not None:
         loaded_batch = Batch(**json_body["batch"])
 
@@ -215,11 +209,8 @@ def save_batches(
             raise Exception('save_batch() "batches" expected to be a set Batch instances')
 
     payload = {"assayId": assay_id, "batches": json_batches}
-    headers = {"Content-type": "application/json", "Accept": "text/plain"}
+    json_body = server_context.make_request(save_batch_url, json=payload)
 
-    json_body = server_context.make_request(
-        save_batch_url, json_dumps(payload, sort_keys=True), headers=headers
-    )
     if json_body is not None:
         resp_batches = json_body["batches"]
         return [Batch(**resp_batch) for resp_batch in resp_batches]

labkey/query.py

@@ -44,9 +44,6 @@
 from typing import List
 
 from .server_context import ServerContext
-from labkey.utils import json_dumps
-
-_query_headers = {"Content-Type": "application/json"}
 
 _default_timeout = 60 * 5  # 5 minutes
 
@@ -176,8 +173,7 @@ def delete_rows(
 
     return server_context.make_request(
         url,
-        json_dumps(payload, sort_keys=True),
-        headers=_query_headers,
+        json=payload,
         timeout=timeout,
     )
 
@@ -203,8 +199,7 @@ def truncate_table(
 
     return server_context.make_request(
         url,
-        json_dumps(payload, sort_keys=True),
-        headers=_query_headers,
+        json=payload,
         timeout=timeout,
     )
 
@@ -294,8 +289,7 @@ def insert_rows(
 
     return server_context.make_request(
         url,
-        json_dumps(payload, sort_keys=True),
-        headers=_query_headers,
+        json=payload,
         timeout=timeout,
     )
 
@@ -430,8 +424,7 @@ def update_rows(
 
     return server_context.make_request(
         url,
-        json_dumps(payload, sort_keys=True),
-        headers=_query_headers,
+        json=payload,
         timeout=timeout,
     )
 

labkey/security.py

@@ -14,12 +14,15 @@
 # limitations under the License.
 #
 import functools
+from dataclasses import dataclass
+
 from typing import Union, List
 
 from labkey.server_context import ServerContext
 
 SECURITY_CONTROLLER = "security"
 USER_CONTROLLER = "user"
+LOGIN_CONTROLLER = "login"
 
 
 def activate_users(
@@ -156,7 +159,7 @@ def get_roles(server_context: ServerContext, container_path: str = None):
     url = server_context.build_url(
         SECURITY_CONTROLLER, "getRoles.api", container_path=container_path
     )
-    return server_context.make_request(url, None)
+    return server_context.make_request(url)
 
 
 def get_user_by_email(server_context: ServerContext, email: str):
@@ -167,7 +170,7 @@ def get_user_by_email(server_context: ServerContext, email: str):
     :return:
     """
     url = server_context.build_url(USER_CONTROLLER, "getUsers.api")
-    payload = dict(includeDeactivatedAccounts=True)
+    payload = {"includeDeactivatedAccounts": True}
     result = server_context.make_request(url, payload)
 
     if result is None or result["users"] is None:
@@ -184,7 +187,6 @@ def list_groups(
     server_context: ServerContext, include_site_groups: bool = False, container_path: str = None
 ):
     url = server_context.build_url(SECURITY_CONTROLLER, "listProjectGroups.api", container_path)
-
     return server_context.make_request(url, {"includeSiteGroups": include_site_groups})
 
 
@@ -242,10 +244,68 @@ def reset_password(server_context: ServerContext, email: str, container_path: st
     :return:
     """
     url = server_context.build_url(SECURITY_CONTROLLER, "adminRotatePassword.api", container_path)
-
     return server_context.make_request(url, {"email": email})
 
 
+def impersonate_user(
+    server_context: ServerContext,
+    user_id: int = None,
+    email: str = None,
+    container_path: str = None,
+):
+    """
+    For site-admins or project-admins only, start impersonating a user.
+
+    Admins may impersonate other users to perform actions on their behalf.
+    Site users may impersonate any user in any project. Project admins must
+    execute this command in a project in which they have admin permission
+    and may impersonate any user that has access to the project.
+
+    To finish an impersonation session use `stop_impersonating`.
+
+    :param user_id: to impersonate (must supply this or email)
+    :param email: to impersonate (must supply this or user_id)
+    :param container_path: in which to impersonate the user
+    """
+    if email is None and user_id is None:
+        raise ValueError("Must supply either [email] or [user_id]")
+
+    url = server_context.build_url(USER_CONTROLLER, "impersonateUser.api", container_path)
+    return server_context.make_request(url, {"userId": user_id, "email": email})
+
+
+def stop_impersonating(server_context: ServerContext):
+    """
+    Stop impersonating a user while keeping the original user logged in.
+    """
+    url = server_context.build_url(LOGIN_CONTROLLER, "stopImpersonating.api")
+    return server_context.make_request(url)
+
+
+@dataclass
+class WhoAmI:
+    id: int
+    email: str
+    display_name: str
+    impersonated: str
+    CSRF: str
+
+
+def who_am_i(server_context: ServerContext) -> WhoAmI:
+    """
+    Calls the whoami API and returns a WhoAmI object.
+    """
+    url = server_context.build_url("login", "whoami.api")
+    response = server_context.make_request(url)
+    return WhoAmI(
+        response["id"],
+        response["email"],
+        response["displayName"],
+        response["impersonated"],
+        response["CSRF"],
+    )
+
+
 def __make_security_group_api_request(
     server_context: ServerContext,
     api: str,
@@ -378,3 +438,15 @@ def remove_from_role(
     @functools.wraps(reset_password)
     def reset_password(self, email: str, container_path: str = None):
         return reset_password(self.server_context, email, container_path)
+
+    @functools.wraps(impersonate_user)
+    def impersonate_user(self, user_id: int = None, email: str = None, container_path: str = None):
+        return impersonate_user(self.server_context, user_id, email, container_path)
+
+    @functools.wraps(stop_impersonating)
+    def stop_impersonating(self):
+        return stop_impersonating(self.server_context)
+
+    @functools.wraps(who_am_i)
+    def who_am_i(self):
+        return who_am_i(self.server_context)

labkey/server_context.py

@@ -1,3 +1,4 @@
+from labkey.utils import json_dumps
 import requests
 from requests.exceptions import RequestException
 from labkey.exceptions import (
@@ -106,44 +107,50 @@ def handle_request_exception(self, exception):
     def make_request(
         self,
         url: str,
-        payload: any,
+        payload: any = None,
         headers: dict = None,
         timeout: int = 300,
         method: str = "POST",
         non_json_response: bool = False,
         file_payload: any = None,
+        json: dict = None,
     ) -> any:
         if self._api_key is not None:
             if self._session.headers.get(API_KEY_TOKEN) is not self._api_key:
                 self._session.headers.update({API_KEY_TOKEN: self._api_key})
 
-        if not self._disable_csrf:
-            if CSRF_TOKEN not in self._session.headers.keys():
-                try:
-                    csrf_url = self.build_url("login", "whoami.api")
-                    response = handle_response(self._session.get(csrf_url))
-                    self._session.headers.update({CSRF_TOKEN: response["CSRF"]})
-                except RequestException as e:
-                    self.handle_request_exception(e)
+        if not self._disable_csrf and CSRF_TOKEN not in self._session.headers.keys():
+            try:
+                csrf_url = self.build_url("login", "whoami.api")
+                response = handle_response(self._session.get(csrf_url))
+                self._session.headers.update({CSRF_TOKEN: response["CSRF"]})
+            except RequestException as e:
+                self.handle_request_exception(e)
 
         try:
             if method == "GET":
-                raw_response = self._session.get(
-                    url, params=payload, headers=headers, timeout=timeout
-                )
+                response = self._session.get(url, params=payload, headers=headers, timeout=timeout)
             else:
                 if file_payload is not None:
-                    raw_response = self._session.post(
+                    response = self._session.post(
                         url,
                         data=payload,
                         files=file_payload,
                         headers=headers,
                         timeout=timeout,
                     )
+                elif json is not None:
+                    if headers is None:
+                        headers = {}
+
+                    headers_ = {**headers, "Content-Type": "application/json"}
+                    # sort_keys is a hack to make unit tests work
+                    data = json_dumps(json, sort_keys=True)
+                    response = self._session.post(url, data=data, headers=headers_, timeout=timeout)
                 else:
-                    raw_response = self._session.post(
+                    response = self._session.post(
                         url, data=payload, headers=headers, timeout=timeout
                     )
-            return handle_response(raw_response, non_json_response)
+            return handle_response(response, non_json_response)
         except RequestException as e:
             self.handle_request_exception(e)

test/integration/conftest.py

@@ -17,9 +17,9 @@
 @pytest.fixture(scope="session")
 def server_context_vars():
     properties_file_path = os.getenv("TEAMCITY_BUILD_PROPERTIES_FILE")
-    host = DEFAULT_HOST
-    port = DEFAULT_PORT
-    context_path = DEFAULT_CONTEXT_PATH
+    host = os.getenv("HOST", DEFAULT_HOST)
+    port = os.getenv("PORT", DEFAULT_PORT)
+    context_path = os.getenv("CONTEXT_PATH", DEFAULT_CONTEXT_PATH)
 
     if properties_file_path is not None:
         with open(properties_file_path) as f:
@@ -45,7 +45,9 @@ def server_context_vars():
 @pytest.fixture(scope="session")
 def api(server_context_vars):
     server, context_path = server_context_vars
-    return APIWrapper(server, PROJECT_NAME, context_path, use_ssl=False)
+    api = APIWrapper(server, PROJECT_NAME, context_path, use_ssl=False)
+    api.security.stop_impersonating()  # Call stop impersonating incase previous test run failed while impersonating a user.
+    return api
 
 
 @pytest.fixture(autouse=True, scope="session")

test/integration/test_security.py

@@ -0,0 +1,47 @@
+from labkey.security import who_am_i
+import pytest
+from labkey.api_wrapper import APIWrapper
+
+# copy from:
+# JavaClientApiTest.testImpersonateUser()
+# JavaClientApiTest.testImpersonationConnection()
+
+pytestmark = pytest.mark.integration  # Mark all tests in this module as integration tests
+TEST_EMAIL = "test_user@test.test"
+TEST_DISPLAY_NAME = "test user"
+
+
+@pytest.fixture(scope="session")
+def test_user(api: APIWrapper, project):
+    url = api.server_context.build_url("security", "createNewUser.api")
+    resp = api.server_context.make_request(url, {"email": TEST_EMAIL, "sendEmail": False})
+    user_id = resp["userId"]
+    yield {"id": user_id, "email": TEST_EMAIL, "display_name": TEST_DISPLAY_NAME}
+    url = api.server_context.build_url("security", "deleteUser.api", container_path="/")
+    resp = api.server_context.make_request(url, {"id": user_id})
+
+
+def test_impersonation(api: APIWrapper, test_user):
+    # test impersonation via email
+    api.security.impersonate_user(email=TEST_EMAIL)
+    who = who_am_i(api.server_context)
+    assert who.display_name == test_user["display_name"]
+    assert who.email == test_user["email"]
+    assert who.id == test_user["id"]
+
+    # test stop impersonating
+    api.security.stop_impersonating()
+    who = who_am_i(api.server_context)
+    assert who.display_name != test_user["display_name"]
+    assert who.email != test_user["email"]
+    assert who.id != test_user["id"]
+
+    # test impersonation via user id
+    api.security.impersonate_user(user_id=test_user["id"])
+    who = who_am_i(api.server_context)
+    assert who.display_name == test_user["display_name"]
+    assert who.email == test_user["email"]
+    assert who.id == test_user["id"]
+
+    # We need to stop impersonating a user before leaving so we don't mess up other tests.
+    api.security.stop_impersonating()

test/unit/test_domain.py

@@ -18,10 +18,7 @@
 import tempfile
 import unittest
 
-try:
-    import mock
-except ImportError:
-    import unittest.mock as mock
+import unittest.mock as mock
 
 from labkey.domain import (
     create,
@@ -70,7 +67,7 @@ class MockCreate(MockLabKey):
 
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
-            "data": json.dumps(domain_definition),
+            "data": json.dumps(domain_definition, sort_keys=True),
             "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
@@ -116,7 +113,7 @@ class MockDrop(MockLabKey):
 
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
-            "data": json.dumps(payload),
+            "data": json.dumps(payload, sort_keys=True),
             "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
@@ -275,7 +272,7 @@ class MockSave(MockLabKey):
 
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
-            "data": json.dumps(payload),
+            "data": json.dumps(payload, sort_keys=True),
             "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
@@ -347,7 +344,7 @@ class MockCreate(MockLabKey):
 
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
-            "data": json.dumps(self.domain_definition),
+            "data": json.dumps(self.domain_definition, sort_keys=True),
             "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
@@ -431,7 +428,7 @@ class MockSave(MockLabKey):
 
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
-            "data": json.dumps(payload),
+            "data": json.dumps(payload, sort_keys=True),
             "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }

test/unit/test_experiment_api.py

@@ -15,10 +15,7 @@
 #
 import unittest
 
-try:
-    import mock
-except ImportError:
-    import unittest.mock as mock
+import unittest.mock as mock
 
 from labkey.experiment import load_batch, save_batch, Batch, Run
 from labkey.exceptions import (
@@ -209,7 +206,7 @@ def setUp(self):
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
             "data": '{"assayId": 12345, "batchId": 54321}',
-            "headers": {"Content-type": "application/json", "Accept": "text/plain"},
+            "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
 
@@ -291,7 +288,7 @@ def setUp(self):
         self.expected_kwargs = {
             "expected_args": [self.service.get_server_url()],
             "data": '{"assayId": 12345, "batches": [{"batchProtocolId": null, "comment": null, "created": null, "createdBy": null, "modified": null, "modifiedBy": null, "name": null, "properties": {"PropertyName": "Property Value"}, "runs": [{"name": "python upload", "properties": {"RunFieldName": "Run Field Value"}}]}]}',
-            "headers": {"Content-type": "application/json", "Accept": "text/plain"},
+            "headers": {"Content-Type": "application/json"},
             "timeout": 300,
         }
 

test/unit/test_query_api.py

@@ -15,10 +15,7 @@
 #
 import unittest
 
-try:
-    import mock
-except ImportError:
-    import unittest.mock as mock
+import unittest.mock as mock
 
 from labkey.query import (
     delete_rows,

test/unit/test_security.py

@@ -15,10 +15,7 @@
 #
 import unittest
 
-try:
-    import mock
-except ImportError:
-    import unittest.mock as mock
+import unittest.mock as mock
 
 from labkey.security import (
     create_user,

test/unit/test_utils.py

@@ -17,10 +17,7 @@
 
 from labkey.server_context import ServerContext
 
-try:
-    import mock
-except ImportError:
-    import unittest.mock as mock
+import unittest.mock as mock
 
 
 def mock_server_context(mock_action):