def test_modify_workspace_properties_Always_On_Dry_Run_True(mocker):
# validate that the stubber call is not maded when Dry Run is set to True
# send an invalid request using stubber and validate that the does not method throws exception
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
response = {}
expected_params = {
'WorkspaceProperties': {'RunningMode': 'ALWAYS_ON'}
}
client_stubber.add_response('modify_workspace_properties', response, expected_params)
client_stubber.activate()
workspace_id = '123qwer'
new_running_mode = 'ALWAYS_ON'
# check if the method throws exception and validate that the stubber was not called
result = workspace_helper.modify_workspace_properties(workspace_id, new_running_mode)
assert result == '-M-'
client_stubber.deactivate()
def test_terminate_unused_workspace_yes(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_id = "123qwe123qwe"
response = {
'FailedRequests': []
}
expected_params = {
'TerminateWorkspaceRequests': [{
'WorkspaceId': workspace_id
}]
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_response('terminate_workspaces', response, expected_params)
client_stubber.activate()
result = workspace_helper.terminate_unused_workspace(workspace_id)
assert result == 'Yes'
client_stubber.deactivate()
def test_get_last_known_user_connection_timestamp(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
last_known_user_connection_timestamp = datetime.datetime.strptime('2021-08-10 19:35:15.524000+00:00',
'%Y-%m-%d %H:%M:%S.%f+00:00')
workspace_id = "123qwe123qwe"
response = {
'WorkspacesConnectionStatus': [{
'LastKnownUserConnectionTimestamp': last_known_user_connection_timestamp
}]
}
expected_params = {
'WorkspaceIds': [workspace_id]
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_response('describe_workspaces_connection_status', response, expected_params)
client_stubber.activate()
result = workspace_helper.get_last_known_user_connection_timestamp(workspace_id)
assert result == last_known_user_connection_timestamp
client_stubber.deactivate()
def test_get_workspaces(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
directory_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
response = {
'Workspaces': [
{'WorkspaceId': '1234'},
{'WorkspaceId': '1234'},
{'WorkspaceId': '1234'}
]
}
expected_params = {
'DirectoryId': directory_id
}
client_stubber.add_response('describe_workspaces', response, expected_params)
client_stubber.activate()
result = workspace_helper.get_workspaces_for_directory(directory_id)
assert result == [
{'WorkspaceId': '1234'},
{'WorkspaceId': '1234'},
{'WorkspaceId': '1234'}
]
client_stubber.deactivate()
def test_get_termination_status_5(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2,
'TerminateUnusedWorkspaces': 'Dry Run'
}
workspace_id = "123qwe123qwe"
billable_time = 1
tags = []
workspace_helper = WorkspacesHelper(settings)
ecs.workspaces_helper.TERMINATE_UNUSED_WORKSPACES = 'Yes'
mocker.patch.object(workspace_helper, 'get_last_known_user_connection_timestamp')
mocker.patch.object(workspace_helper, 'check_workspace_usage_for_current_month')
workspace_helper.check_workspace_usage_for_current_month.return_value = False
mocker.patch.object(workspace_helper, 'check_if_workspace_available_on_first_day')
workspace_helper.check_if_workspace_available_on_first_day.return_value = True
mocker.patch.object(workspace_helper, 'check_if_workspace_needs_to_be_terminated')
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_check_workspace_usage_for_current_month_None(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
last_known_user_connection_timestamp = None
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.check_workspace_usage_for_current_month(last_known_user_connection_timestamp)
assert result is True
def test_check_workspace_usage_for_current_month_true(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
last_known_user_connection_timestamp = datetime.datetime.utcnow().today().replace(day=1, hour=0, minute=0, second=0,
microsecond=0)
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.check_workspace_usage_for_current_month(last_known_user_connection_timestamp)
assert result is True
def test_check_if_workspace_needs_to_be_terminated_3(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2,
'terminateUnusedWorkspaces': 'Yes'
}
workspace_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.check_if_workspace_needs_to_be_terminated(workspace_id)
assert result == ''
def test_check_for_skip_tag_false(mocker):
tags = [{'Key': 'nothing', 'Value': 'True'}]
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': 'yes',
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.check_for_skip_tag(tags)
assert result is False
def test_check_workspace_usage_for_current_month_false(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
last_known_user_connection_timestamp = datetime.datetime.strptime('2021-01-10 19:35:15.524000+00:00',
'%Y-%m-%d %H:%M:%S.%f+00:00')
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.check_workspace_usage_for_current_month(last_known_user_connection_timestamp)
assert result is False
def test_check_if_workspace_needs_to_be_terminated_5(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': False,
'startTime': 1,
'endTime': 2,
'terminateUnusedWorkspaces': 'Yes'
}
workspace_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
mocker.patch.object(workspace_helper, 'terminate_unused_workspace')
workspace_helper.terminate_unused_workspace.return_value = ''
result = workspace_helper.check_if_workspace_needs_to_be_terminated(workspace_id)
assert result == ''
def test_process_workspace_performance(mocker):
workspace = {
'WorkspaceId': 'ws-68h123hty',
'DirectoryId': 'd-901230bb84',
'UserName': '******',
'IpAddress': '111.16.1.233',
'State': 'AVAILABLE',
'BundleId': 'wsb-cl123qzj1',
'SubnetId': 'subnet-05d421387eaa7cf86',
'ComputerName': 'A-APPW123KP4NP',
'WorkspaceProperties': {
'RunningMode': 'ALWAYS_ON',
'RootVolumeSizeGib': 80,
'UserVolumeSizeGib': 50,
'ComputeTypeName': 'PERFORMANCE'
},
'ModificationStates': []
}
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
mocker.patch.object(workspace_helper.metrics_helper, 'get_billable_hours')
workspace_helper.metrics_helper.get_billable_hours.return_value = 100
mocker.patch.object(workspace_helper, 'get_tags')
mocker.patch.object(workspace_helper, 'check_for_skip_tag')
workspace_helper.check_for_skip_tag.return_value = False
mocker.patch.object(workspace_helper, 'get_hourly_threshold')
workspace_helper.get_hourly_threshold.return_value = 5
mocker.patch.object(workspace_helper, 'compare_usage_metrics')
workspace_helper.compare_usage_metrics.return_value = {
'resultCode': '-N-',
'newMode': 'ALWAYS_ON'
}
mocker.patch.object(workspace_helper, 'get_termination_status')
result = workspace_helper.process_workspace(workspace)
assert result['bundleType'] == 'PERFORMANCE'
assert result['billableTime'] == 100
def test_get_workspaces_exception(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
directory_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_client_error('describe_workspaces', 'Invalid Directory')
client_stubber.activate()
result = workspace_helper.get_workspaces_for_directory(directory_id)
assert result == []
client_stubber.deactivate()
def test_check_workspace_available_exception(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.cloudwatch_client)
workspace_id = '123qwer'
client_stubber.add_client_error('get_metric_statistics', 'Invalid request')
client_stubber.activate()
result = workspace_helper.check_if_workspace_available_on_first_day(workspace_id)
assert result is False
client_stubber.deactivate()
def test_terminate_unused_workspace_exception(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_client_error('terminate_workspaces', 'Invalid Directory')
client_stubber.activate()
result = workspace_helper.terminate_unused_workspace(workspace_id)
assert result == ''
client_stubber.deactivate()
def test_get_termination_status_1():
ecs.workspaces_helper.TERMINATE_UNUSED_WORKSPACES = 'No'
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2,
'TerminateUnusedWorkspaces': 'Yes'
}
workspace_id = "123qwe123qwe"
billable_time = 0
tags = []
workspace_helper = WorkspacesHelper(settings)
result = workspace_helper.get_termination_status(workspace_id, billable_time, tags)
assert result == ''
def test_get_last_known_user_connection_timestamp_exception(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_id = "123qwe123qwe"
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_client_error('get_last_known_user_connection_timestamp', 'Invalid workspace')
client_stubber.activate()
result = workspace_helper.get_last_known_user_connection_timestamp(workspace_id)
assert result is None
client_stubber.deactivate()
def test_terminate_workspaces_no(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
client_stubber.add_client_error('TerminateWorkspaceRequests', "Invalid_request")
client_stubber.activate()
workspace_id = '123qwer'
result = workspace_helper.terminate_unused_workspace(workspace_id)
assert result == ''
client_stubber.deactivate()
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 == []
def test_check_workspace_available_yes(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': True,
'startTime': 1,
'endTime': 2
}
start_time = time.strftime("%Y-%m") + '-01T00:00:00Z'
end_time = time.strftime("%Y-%m") + '-02T00:00:00Z'
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.cloudwatch_client)
workspace_id = '123qwer'
response = {
'Datapoints': [
{'Timestamp': datetime.datetime(2021, 5, 2, 11, 0, tzinfo=tzutc()), 'Maximum': 1.0, 'Unit': 'Count'},
{'Timestamp': datetime.datetime(2021, 5, 1, 7, 0, tzinfo=tzutc()), 'Maximum': 1.0, 'Unit': 'Count'},
{'Timestamp': datetime.datetime(2021, 5, 2, 6, 0, tzinfo=tzutc()), 'Maximum': 1.0, 'Unit': 'Count'},
{'Timestamp': datetime.datetime(2021, 5, 1, 2, 0, tzinfo=tzutc()), 'Maximum': 0.0, 'Unit': 'Count'},
{'Timestamp': datetime.datetime(2021, 5, 2, 1, 0, tzinfo=tzutc()), 'Maximum': 0.0, 'Unit': 'Count'}
]
}
expected_params = {
'Dimensions': [
{
'Name': 'WorkspaceId',
'Value': workspace_id
}
],
'Namespace': 'AWS/WorkSpaces',
'MetricName': 'Available',
'StartTime': start_time,
'EndTime': end_time,
'Period': 300,
'Statistics': ['Maximum']
}
client_stubber.add_response('get_metric_statistics', response, expected_params)
client_stubber.activate()
result = workspace_helper.check_if_workspace_available_on_first_day(workspace_id)
assert result is True
client_stubber.deactivate()
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_modify_workspace_properties_Always_On(mocker):
settings = {
'region': 'us-east-1',
'hourlyLimits': 10,
'testEndOfMonth': 'yes',
'isDryRun': False,
'startTime': 1,
'endTime': 2
}
workspace_helper = WorkspacesHelper(settings)
client_stubber = Stubber(workspace_helper.workspaces_client)
response = {}
expected_params = {
'WorkspaceId': '123qwer',
'WorkspaceProperties': {'RunningMode': 'ALWAYS_ON'}
}
client_stubber.add_response('modify_workspace_properties', response, expected_params)
client_stubber.activate()
workspace_id = '123qwer'
new_running_mode = 'ALWAYS_ON'
result = workspace_helper.modify_workspace_properties(workspace_id, new_running_mode)
assert result == '-M-'
def process_directory(self, region, stack_parameters,
directory_parameters):
workspace_count = 0
end_time = directory_parameters['EndTime']
start_time = directory_parameters['StartTime']
list_processed_workspaces = []
directory_csv = ''
log_body_directory_csv = ''
is_dry_run = self.get_dry_run(stack_parameters)
test_end_of_month = self.get_end_of_month(stack_parameters)
directory_id = directory_parameters['DirectoryId']
report_csv = 'WorkspaceID,Billable Hours,Usage Threshold,Change Reported,Bundle Type,Initial Mode,New Mode,Username,Computer Name,DirectoryId,WorkspaceTerminated,Tags\n'
# List of bundles with specific hourly limits
workspaces_helper = WorkspacesHelper({
'region':
region,
'hourlyLimits': {
'VALUE': stack_parameters['ValueLimit'],
'STANDARD': stack_parameters['StandardLimit'],
'PERFORMANCE': stack_parameters['PerformanceLimit'],
'POWER': stack_parameters['PowerLimit'],
'POWERPRO': stack_parameters['PowerProLimit'],
'GRAPHICS': stack_parameters['GraphicsLimit'],
'GRAPHICSPRO': stack_parameters['GraphicsProLimit']
},
'testEndOfMonth':
test_end_of_month,
'isDryRun':
is_dry_run,
'startTime':
start_time,
'endTime':
end_time,
'terminateUnusedWorkspaces':
stack_parameters['TerminateUnusedWorkspaces']
})
list_workspaces = workspaces_helper.get_workspaces_for_directory(
directory_id)
for workspace in list_workspaces:
log.debug("Processing workspace {}".format(workspace))
workspace_count = workspace_count + 1
result = workspaces_helper.process_workspace(workspace)
report_csv = workspaces_helper.append_entry(
report_csv, result) # Append result data to the CSV
directory_csv = workspaces_helper.append_entry(
directory_csv, result) # Append result for aggregated report
try:
workspace_processed = {
'previousMode': result['initialMode'],
'newMode': result['newMode'],
'bundleType': result['bundleType'],
'hourlyThreshold': result['hourlyThreshold'],
'billableTime': result['billableTime']
}
list_processed_workspaces.append(workspace_processed)
except Exception:
log.debug(
"Could not append workspace for metrics. Skipping this workspace"
)
log_body = workspaces_helper.expand_csv(report_csv)
log_body_directory_csv = workspaces_helper.expand_csv(
directory_csv)
upload_report(stack_parameters, log_body, directory_id, region)
return workspace_count, list_processed_workspaces, log_body_directory_csv