fix legal hold/API client bug in devices command (#396)

* fix bug in `devices list --include-legal-hold-membership` when using api client auth

* black

* skip new bugbear warning

* version and changelog

Author: timabrmsn

Diff for: CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 The intended audience of this file is for py42 consumers -- as such, changes that don't affect
 how a consumer would use the library (e.g. adding unit tests, updating documentation, etc) are not captured here.
 
+## 1.16.3 - 2022-12-08
+
+### Fixed
+
+- Bug in `devices list` command when using `--include-legal-hold-membership` option with an API client auth profile.
+
 ## 1.16.2 - 2022-11-07
 
 ### Fixed

Diff for: src/code42cli/__version__.py

@@ -1 +1 @@
-__version__ = "1.16.2"
+__version__ = "1.16.3"

Diff for: src/code42cli/cmds/devices.py

@@ -432,11 +432,21 @@ def _add_legal_hold_membership_to_device_dataframe(sdk, df):
 
 def _get_all_active_hold_memberships(sdk):
     for page in sdk.legalhold.get_all_matters(active=True):
-        for matter in page["legalHolds"]:
-            for _page in sdk.legalhold.get_all_matter_custodians(
-                legal_hold_uid=matter["legalHoldUid"], active=True
-            ):
-                yield from _page["legalHoldMemberships"]
+        if sdk._auth_flag == 1:  # noqa: api client endpoint returns a list directly
+            matters = page.data
+        else:
+            matters = page["legalHolds"]
+        for matter in matters:
+            if sdk._auth_flag == 1:  # noqa: api client endpoint returns a list directly
+                for _page in sdk.legalhold.get_all_matter_custodians(
+                    legal_hold_matter_uid=matter["legalHoldUid"], active=True
+                ):
+                    yield from _page.data
+            else:
+                for _page in sdk.legalhold.get_all_matter_custodians(
+                    legal_hold_uid=matter["legalHoldUid"], active=True
+                ):
+                    yield from _page["legalHoldMemberships"]
 
 
 def _get_device_dataframe(

Diff for: src/code42cli/output_formats.py

@@ -138,7 +138,7 @@ def _iter_json(self, dfs, columns=None, **kwargs):
             yield f"{json_string}\n"
 
     def _checkpoint_and_iter_formatted_events(self, df, formatted_rows):
-        for event, row in zip(df.to_dict("records"), formatted_rows):
+        for event, row in zip(df.to_dict("records"), formatted_rows):  # noqa: B905
             yield row
             self.checkpoint_func(event)
 
@@ -188,7 +188,7 @@ def iter_rows(self, dfs, columns=None):
                 filtered = self._select_columns(df, columns)
             else:
                 filtered = df
-            for full_event, filtered_event in zip(
+            for full_event, filtered_event in zip(  # noqa: B905
                 df.to_dict("records"), filtered.to_dict("records")
             ):
                 yield filtered_event

Diff for: tests/cmds/test_devices.py

@@ -279,6 +279,42 @@
         },
     ]
 }
+API_CLIENT_MATTER_RESPONSE = [
+    {
+        "legalHoldUid": "123456789",
+        "name": "Test legal hold matter",
+        "description": "",
+        "notes": None,
+        "holdExtRef": None,
+        "active": True,
+        "creationDate": "2020-08-05T10:49:58.353-05:00",
+        "lastModified": "2020-08-05T10:49:58.358-05:00",
+        "creator": {
+            "userUid": "12345",
+            "username": "[email protected]",
+            "email": "[email protected]",
+            "userExtRef": None,
+        },
+        "holdPolicyUid": "966191295667423997",
+    },
+    {
+        "legalHoldUid": "987654321",
+        "name": "Another Matter",
+        "description": "",
+        "notes": None,
+        "holdExtRef": None,
+        "active": True,
+        "creationDate": "2020-05-20T15:58:31.375-05:00",
+        "lastModified": "2020-05-28T13:49:16.098-05:00",
+        "creator": {
+            "userUid": "76543",
+            "username": "[email protected]",
+            "email": "[email protected]",
+            "userExtRef": None,
+        },
+        "holdPolicyUid": "946178665645035826",
+    },
+]
 ALL_CUSTODIANS_RESPONSE = {
     "legalHoldMemberships": [
         {
@@ -310,6 +346,35 @@
         },
     ]
 }
+API_CLIENT_ALL_CUSTODIANS_RESPONSE = [
+    {
+        "legalHoldMembershipUid": "99999",
+        "active": True,
+        "creationDate": "2020-07-16T08:50:23.405Z",
+        "legalHold": {
+            "legalHoldUid": "123456789",
+            "name": "Test legal hold matter",
+        },
+        "user": {
+            "userUid": "840103986007089121",
+            "username": "[email protected]",
+            "email": "[email protected]",
+            "userExtRef": None,
+        },
+    },
+    {
+        "legalHoldMembershipUid": "88888",
+        "active": True,
+        "creationDate": "2020-07-16T08:50:23.405Z",
+        "legalHold": {"legalHoldUid": "987654321", "name": "Another Matter"},
+        "user": {
+            "userUid": "840103986007089121",
+            "username": "[email protected]",
+            "email": "[email protected]",
+            "userExtRef": None,
+        },
+    },
+]
 
 
 @pytest.fixture
@@ -355,12 +420,18 @@ def users_list_generator():
     yield TEST_USERS_LIST_PAGE
 
 
-def matter_list_generator():
-    yield MATTER_RESPONSE
+def matter_list_generator(mocker, api_client=False):
+    if api_client:
+        yield create_mock_response(mocker, data=API_CLIENT_MATTER_RESPONSE)
+    else:
+        yield create_mock_response(mocker, data=MATTER_RESPONSE)
 
 
-def custodian_list_generator():
-    yield ALL_CUSTODIANS_RESPONSE
+def custodian_list_generator(mocker, api_client=False):
+    if api_client:
+        yield create_mock_response(mocker, data=API_CLIENT_ALL_CUSTODIANS_RESPONSE)
+    else:
+        yield create_mock_response(mocker, data=ALL_CUSTODIANS_RESPONSE)
 
 
 @pytest.fixture
@@ -446,14 +517,28 @@ def get_all_users_success(cli_state):
 
 
 @pytest.fixture
-def get_all_matter_success(cli_state):
-    cli_state.sdk.legalhold.get_all_matters.return_value = matter_list_generator()
+def get_all_matter_success(mocker, cli_state):
+    cli_state.sdk.legalhold.get_all_matters.return_value = matter_list_generator(mocker)
+
+
+@pytest.fixture
+def get_api_client_all_matter_success(mocker, cli_state):
+    cli_state.sdk.legalhold.get_all_matters.return_value = matter_list_generator(
+        mocker, api_client=True
+    )
 
 
 @pytest.fixture
-def get_all_custodian_success(cli_state):
+def get_all_custodian_success(mocker, cli_state):
     cli_state.sdk.legalhold.get_all_matter_custodians.return_value = (
-        custodian_list_generator()
+        custodian_list_generator(mocker)
+    )
+
+
+@pytest.fixture
+def get_api_client_all_custodian_success(mocker, cli_state):
+    cli_state.sdk.legalhold.get_all_matter_custodians.return_value = (
+        custodian_list_generator(mocker, api_client=True)
     )
 
 
@@ -740,6 +825,21 @@ def test_add_legal_hold_membership_to_device_dataframe_adds_legal_hold_columns_t
     assert "legalHoldName" in result.columns
 
 
+def test_api_client_add_legal_hold_membership_to_device_dataframe_adds_legal_hold_columns_to_dataframe(
+    cli_state, get_api_client_all_matter_success, get_api_client_all_custodian_success
+):
+    cli_state.sdk._auth_flag = 1
+    testdf = DataFrame.from_records(
+        [
+            {"userUid": "840103986007089121", "status": "Active"},
+            {"userUid": "836473273124890369", "status": "Active, Deauthorized"},
+        ]
+    )
+    result = _add_legal_hold_membership_to_device_dataframe(cli_state.sdk, testdf)
+    assert "legalHoldUid" in result.columns
+    assert "legalHoldName" in result.columns
+
+
 def test_list_without_page_size_option_defaults_to_100_results_per_page(
     cli_state, runner
 ):

Diff for: tests/conftest.py

@@ -297,7 +297,7 @@ def mock_dataframe_to_string(mocker):
 
 
 def create_mock_response(mocker, data=None, status=200):
-    if isinstance(data, dict):
+    if isinstance(data, (dict, list)):
         data = json.dumps(data)
     elif not data:
         data = ""