Update to version 2.4.1

Author: hayesry

.DS_Store

@@ -3,6 +3,10 @@
 
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+ ## [2.4.1] - 2021-10
+ ### Fixed
+ - Fixed the bug to get all the workspaces in a directory
+ 
  ## [2.4] - 2021-09
  ### Fixed
  - Fixed the bug to correctly calculate billable hours if user disconnects workspace within autostop timeout

CHANGELOG.md

@@ -96,6 +96,7 @@ def process_workspace(self, workspace):
             workspace_terminated = self.get_termination_status(workspace_id, billable_time, tags)
             optimization_result = self.compare_usage_metrics(workspace_id, billable_time, hourly_threshold,
                                                              workspace_running_mode)
+
         return {
             'workspaceID': workspace_id,
             'billableTime': billable_time,
@@ -104,9 +105,9 @@ def process_workspace(self, workspace):
             'newMode': optimization_result['newMode'],
             'bundleType': workspace_bundle_type,
             'initialMode': workspace_running_mode,
-            'userName': workspace['UserName'],
-            'computerName': workspace['ComputerName'],
-            'directoryId': workspace['DirectoryId'],
+            'userName': workspace.get('UserName', ''),
+            'computerName': workspace.get('ComputerName', ''),
+            'directoryId': workspace.get('DirectoryId', ''),
             'tags': tags,
             'workspaceTerminated': workspace_terminated
         }
@@ -200,7 +201,7 @@ def get_workspaces_for_directory(self, directory_id):
                     DirectoryId=directory_id,
                     NextToken=next_token
                 )
-                list_workspaces.extend(response.get('Directories', []))
+                list_workspaces.extend(response.get('Workspaces', []))
                 next_token = response.get('NextToken', None)
         except botocore.exceptions.ClientError as e:
             log.error(

source/.coverage

@@ -878,3 +878,102 @@ def test_get_termination_status_5(mocker):
     workspace_helper.check_if_workspace_needs_to_be_terminated.return_value = 'Yes - Dry Run'
     result = workspace_helper.get_termination_status(workspace_id, billable_time, tags)
     assert result == ''
+
+
+def test_get_workspaces_for_directory_use_next_token():
+    settings = {
+        'region': 'us-east-1',
+        'hourlyLimits': 10,
+        'testEndOfMonth': 'yes',
+        'isDryRun': True,
+        'startTime': 1,
+        'endTime': 2,
+        'TerminateUnusedWorkspaces': 'Dry Run'
+    }
+    directory_id = "123qwe123qwe"
+    workspace_helper = WorkspacesHelper(settings)
+    client_stubber = Stubber(workspace_helper.workspaces_client)
+
+    expected_params_1 = {
+        'DirectoryId': directory_id
+    }
+
+    response_1 = {
+    'Workspaces': [{'WorkspaceId':'id_1'}],
+    'NextToken': 's223123jj32'
+    }
+
+    expected_params_2 = {
+        'DirectoryId': directory_id,
+        'NextToken': 's223123jj32'
+    }
+
+    response_2 = {
+        'Workspaces': [{'WorkspaceId':'id_2'}]
+    }
+
+    client_stubber.add_response('describe_workspaces', response_1, expected_params_1)
+    client_stubber.add_response('describe_workspaces', response_2, expected_params_2)
+    client_stubber.activate()
+    response = workspace_helper.get_workspaces_for_directory(directory_id)
+    client_stubber.activate()
+    assert response == [{'WorkspaceId': 'id_1'}, {'WorkspaceId': 'id_2'}]
+
+
+def test_get_workspaces_for_directory_no_next_token():
+    settings = {
+        'region': 'us-east-1',
+        'hourlyLimits': 10,
+        'testEndOfMonth': 'yes',
+        'isDryRun': True,
+        'startTime': 1,
+        'endTime': 2,
+        'TerminateUnusedWorkspaces': 'Dry Run'
+    }
+    directory_id = "123qwe123qwe"
+    workspace_helper = WorkspacesHelper(settings)
+    client_stubber = Stubber(workspace_helper.workspaces_client)
+
+    expected_params_1 = {
+        'DirectoryId': directory_id
+    }
+
+    response_1 = {
+    'Workspaces': [{'WorkspaceId':'id_1'}]
+    }
+
+    expected_params_2 = {
+        'DirectoryId': directory_id,
+        'NextToken': 's223123jj32'
+    }
+
+    response_2 = {
+        'Workspaces': [{'WorkspaceId':'id_2'}]
+    }
+
+    client_stubber.add_response('describe_workspaces', response_1, expected_params_1)
+    client_stubber.add_response('describe_workspaces', response_2, expected_params_2)
+    client_stubber.activate()
+    response = workspace_helper.get_workspaces_for_directory(directory_id)
+    client_stubber.activate()
+    assert response == [{'WorkspaceId': 'id_1'}]
+
+
+def test_get_workspaces_for_directory_return_exception():
+    settings = {
+        'region': 'us-east-1',
+        'hourlyLimits': 10,
+        'testEndOfMonth': 'yes',
+        'isDryRun': True,
+        'startTime': 1,
+        'endTime': 2,
+        'TerminateUnusedWorkspaces': 'Dry Run'
+    }
+    directory_id = "123qwe123qwe"
+    workspace_helper = WorkspacesHelper(settings)
+    client_stubber = Stubber(workspace_helper.workspaces_client)
+    client_stubber.add_client_error('describe_workspaces', "Invalid_request")
+    client_stubber.activate()
+    response = workspace_helper.get_workspaces_for_directory(directory_id)
+    client_stubber.activate()
+    assert response == []