"""Tests of disco_ssm"""
import random
import copy
import json
from unittest import TestCase
from mock import MagicMock, patch, call

from botocore.exceptions import ClientError

from disco_aws_automation import DiscoSSM
from disco_aws_automation import disco_ssm
from disco_aws_automation.disco_ssm import SSM_DOCUMENTS_DIR, SSM_OUTPUT_ERROR_DELIMITER

from tests.helpers.patch_disco_aws import (get_mock_config,
                                           TEST_ENV_NAME)

MOCK_BUCKET_PREFIX = 'asiaq-test-buckets'
MOCK_S3_BUCKET_NAME = MOCK_BUCKET_PREFIX + '--ssm'

MOCK_AWS_CONFIG_DEFINITION = {
    'disco_aws': {
        'default_environment': TEST_ENV_NAME,
        's3_bucket_base': MOCK_BUCKET_PREFIX,
    }}
MOCK_AWS_DOCUMENTS = [
    {
        'Name': 'AWS-document_1',
        'Owner': 'mock_owner',
        'PlatformTypes': ['Linux']
    },
    {
        'Name': 'AWS-document_2',
        'Owner': 'mock_owner',
        'PlatformTypes': ['Linux']
    }
]
MOCK_ASIAQ_DOCUMENTS = [
    {
        'Name': 'asiaq-ssm_document_1',
        'Owner': 'mock_owner',
        'PlatformTypes': ['Linux']
    },
    {
        'Name': 'asiaq-ssm_document_2',
        'Owner': 'mock_owner',
        'PlatformTypes': ['Linux']
    }
]
MOCK_ASIAQ_DOCUMENT_CONTENTS = {
    'asiaq-ssm_document_1': '{"field1": "value1"}',
    'asiaq-ssm_document_2': '{"field2": "value2"}',
}
MOCK_ASIAQ_DOCUMENT_FILE_CONTENTS = {
    SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_1.ssm': MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_1'],
    SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_2.ssm': MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_2'],
}
MOCK_NEXT_TOKEN = "abcdefgABCDEFG"
MOCK_COMMANDS = []
MOCK_COMMAND_INVOCATIONS = []
MOCK_S3_STORE = {}


# pylint: disable=invalid-name
def mock_boto3_client(arg):
    """ mock method for boto3.client() """
    if arg == "ssm":
        return _get_mock_ssm()
    elif arg == "s3":
        return _get_mock_s3()
    else:
        raise Exception("Mock {} client not implemented.".format(arg))


def _get_mock_ssm():
    """Get mock ssm client"""
    mock_asiaq_documents = copy.copy(MOCK_ASIAQ_DOCUMENTS)
    mock_asiaq_document_contents = copy.copy(MOCK_ASIAQ_DOCUMENT_CONTENTS)
    wait_flags = {'delete': True, 'create': True}

    def _mock_list_documents(NextToken=''):
        all_documents = MOCK_AWS_DOCUMENTS + mock_asiaq_documents
        if NextToken == '':
            return {
                'DocumentIdentifiers': all_documents[:len(all_documents) / 2],
                'NextToken': MOCK_NEXT_TOKEN
            }
        elif NextToken == MOCK_NEXT_TOKEN:
            return {
                'DocumentIdentifiers': all_documents[len(all_documents) / 2:],
                'NextToken': ''
            }
        else:
            raise RuntimeError("Invalid NextToken: {0}".format(NextToken))

    def _mock_get_document(Name):
        if Name not in mock_asiaq_document_contents:
            raise ClientError({'Error': {'Code': 'Mock_code', 'Message': 'mock message'}},
                              'GetDocument')

        return {'Name': Name,
                'Content': mock_asiaq_document_contents[Name]}

    def _mock_create_document(Content, Name):
        mock_asiaq_documents.append({'Name': Name,
                                     'Owner': 'mock_owner',
                                     'PlatformTypes': ['Linux']})
        mock_asiaq_document_contents[Name] = Content

    def _mock_delete_document(Name):
        doc_to_delete = [document for document in mock_asiaq_documents
                         if document['Name'] == Name]
        if doc_to_delete:
            mock_asiaq_documents.remove(doc_to_delete[0])
        mock_asiaq_document_contents.pop(Name, None)

    def _mock_describe_document(Name):
        # Using two wait flags to simulate that AWS is taking time to delete and
        # create documents
        if Name not in mock_asiaq_document_contents:
            if wait_flags['delete']:
                wait_flags['delete'] = False
                return {'Document': {'Name': Name, 'Status': 'Active'}}
            else:
                wait_flags['delete'] = True
                raise ClientError({'Error': {'Code': 'Mock_code', 'Message': 'mock message'}},
                                  'DescribeDocument')
        else:
            if wait_flags['create']:
                wait_flags['create'] = False
                res = {'Document': {'Name': Name, 'Status': 'Creating'}}
            else:
                wait_flags['create'] = True
                res = {'Document': {'Name': Name, 'Status': 'Active'}}
            return res

    def _mock_send_command(InstanceIds, DocumentName, Comment=None, Parameters=None, OutputS3BucketName=None):
        mock_command = _create_mock_command(
            InstanceIds,
            DocumentName,
            Comment,
            Parameters,
            OutputS3BucketName
        )
        return {"Command": mock_command}

    def _mock_list_commands(CommandId):
        filtered_mock_commands = []

        for command in MOCK_COMMANDS:
            if command['CommandId'] == CommandId:
                filtered_mock_commands.append(command)
                break

        return {"Commands": filtered_mock_commands}

    def _mock_list_command_invocations(CommandId, Details):
        if not Details:
            raise Exception("This mock method always expects details, so not wanting them is unsupported.")

        filtered_mock_command_invocations = []

        for command_invocation in MOCK_COMMAND_INVOCATIONS:
            if command_invocation['CommandId'] == CommandId:
                filtered_mock_command_invocations.append(command_invocation)

        return {"CommandInvocations": filtered_mock_command_invocations}

    mock_ssm = MagicMock()
    mock_ssm.list_documents.side_effect = _mock_list_documents
    mock_ssm.get_document.side_effect = _mock_get_document
    mock_ssm.create_document.side_effect = _mock_create_document
    mock_ssm.delete_document.side_effect = _mock_delete_document
    mock_ssm.describe_document.side_effect = _mock_describe_document
    mock_ssm.send_command.side_effect = _mock_send_command
    mock_ssm.list_commands.side_effect = _mock_list_commands
    mock_ssm.list_command_invocations.side_effect = _mock_list_command_invocations

    return mock_ssm


def _get_mock_s3():
    """Get mock s3 client"""
    def _mock_list_objects_v2(Bucket, Prefix):
        data = MOCK_S3_STORE[Bucket]

        if data.keys():
            return {
                'Contents': [{'Key': key} for key in data.keys() if key.startswith(Prefix)]
            }

        return {}

    def _mock_head_bucket(Bucket):
        if Bucket != MOCK_S3_BUCKET_NAME:
            raise ClientError({'Error': {}}, 'HeadBucket')
        return None

    def _mock_get_object(Bucket, Key):
        value = MOCK_S3_STORE[Bucket][Key]

        body = MagicMock()
        body.read.return_value = value

        response = {'Body': body}

        return response

    mock_s3 = MagicMock()
    mock_s3.list_objects_v2.side_effect = _mock_list_objects_v2
    mock_s3.head_bucket.side_effect = _mock_head_bucket
    mock_s3.get_object.side_effect = _mock_get_object

    return mock_s3


def _combine_stdout_and_stderr(stdout=None, stderr=None):
    if stdout is None and stderr is None:
        return ''
    elif stdout is not None and stderr is None:
        return stdout

    return SSM_OUTPUT_ERROR_DELIMITER.join([stdout or '', stderr or ''])


def _create_mock_command(instance_ids, document_name, comment=None, parameters=None,
                         output_s3_bucket_name=None, status='Success', stdout='stdout', stderr='stderr'):
    command_id = ''.join(random.choice("0123456789abcdefghijklmnopqrstuvxyz-") for _ in range(40))

    mock_command = {
        "Status": status,
        "Parameters": parameters or {},
        "DocumentName": document_name,
        "InstanceIds": instance_ids,
        "CommandId": command_id
    }

    if output_s3_bucket_name is not None:
        mock_command["OutputS3BucketName"] = output_s3_bucket_name

    if comment is not None:
        mock_command["Comment"] = comment

    mock_invocations = []

    for instance_id in instance_ids:
        mock_invocation = {
            "Status": status,
            "CommandPlugins": [
                {
                    "Status": status,
                    "Name": 'foo-plugin',
                    "ResponseCode": 0,
                    "Output": _combine_stdout_and_stderr(stdout, stderr)
                }
            ],
            "InstanceId": instance_id,
            "DocumentName": document_name,
            "CommandId": command_id
        }

        if output_s3_bucket_name is not None:
            mock_invocation["CommandPlugins"][0]["OutputS3BucketName"] = output_s3_bucket_name
            mock_invocation["CommandPlugins"][0]["OutputS3KeyPrefix"] = "{0}/{1}".format(
                command_id,
                instance_id
            )

        if comment is not None:
            mock_invocation["Comment"] = comment

        mock_invocations.append(mock_invocation)

    MOCK_COMMANDS.append(mock_command)
    MOCK_COMMAND_INVOCATIONS.extend(mock_invocations)

    if output_s3_bucket_name is not None:
        mock_s3_bucket = create_mock_s3_bucket(mock_invocations, stdout, stderr)
        MOCK_S3_STORE[output_s3_bucket_name] = mock_s3_bucket

    return mock_command


def create_mock_s3_bucket(invocations, stdout, stderr):
    """
    Creates a mock Asiaq S3 Bucket object that contains the output of the given command invocations
    """
    data = {}

    for invocation in invocations:
        prefix = invocation["CommandPlugins"][0]["OutputS3KeyPrefix"]

        if stdout is not None:
            stdout_key = "{0}/{1}".format(prefix, 'stdout')
            data[stdout_key] = stdout

        if stderr is not None:
            stderr_key = "{0}/{1}".format(prefix, 'stderr')
            data[stderr_key] = stderr

    return data


def create_mock_open(content_dict):
    """
    Creates a mock open method that returns the file content based
    on the dict being passed in
    """
    def _mock_open(file_name, mode):
        if file_name not in content_dict:
            raise RuntimeError("File name ({0}) not in content dict.".format(file_name))

        if mode != 'r':
            raise RuntimeError("SSM documents should only be opened by asiaq in read mode.")

        mock_file = MagicMock(spec=file)
        mock_file.__enter__.return_value = mock_file
        mock_file.read.return_value = content_dict[file_name]
        return mock_file

    return _mock_open


def _standardize_json_str(json_str):
    return json.dumps(json.loads(json_str), indent=4)


class DiscoSSMTests(TestCase):
    """Test DiscoSSM"""

    def setUp(self):
        config_aws = get_mock_config(MOCK_AWS_CONFIG_DEFINITION)
        self._ssm = DiscoSSM(environment_name=TEST_ENV_NAME,
                             config_aws=config_aws)

    @patch('boto3.client', mock_boto3_client)
    def test_get_all_documents(self):
        """Verify that get_all_documents() works"""

        # Calling the method under test
        documents = self._ssm.get_all_documents()

        # Begin verifications
        # Make sure the documents returned contain only the asiaq-managed ones
        self.assertEqual(documents, MOCK_ASIAQ_DOCUMENTS)

        expected_list_calls = [call(), call(NextToken=MOCK_NEXT_TOKEN)]
        self.assertEqual(expected_list_calls,
                         self._ssm.conn.list_documents.mock_calls)

    @patch('boto3.client', mock_boto3_client)
    def test_get_document_content(self):
        """Verify content of a document is correctly retrieved"""

        # Calling the method under test
        doc_content_1 = self._ssm.get_document_content('asiaq-ssm_document_1')
        doc_content_2 = self._ssm.get_document_content('asiaq-ssm_document_2')

        # Verify result
        self.assertEqual(doc_content_1, MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_1'])
        self.assertEqual(doc_content_2, MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_2'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_document_invalid_doc_name(self):
        """Verify correct error is thrown when doc_name is invalid"""

        # Calling the method under test
        self.assertRaises(Exception, self._ssm.get_document_content, doc_name='AWS-doc')

    @patch('boto3.client', mock_boto3_client)
    def test_get_document_doc_name_not_found(self):
        """Verify no content is returned when doc_name is not found"""

        # Calling the method under test
        doc_content = self._ssm.get_document_content('asiaq-random_doc')

        # Verifying results
        self.assertEqual(doc_content, None)

    @patch('boto3.client', mock_boto3_client)
    @patch('os.listdir')
    @patch('disco_aws_automation.disco_ssm.open')
    def test_update_create_docs(self, mock_open, mock_os_listdir):
        """Verify that creating new documents in the update() method works"""
        # Setting up test
        mock_os_listdir.return_value = ['asiaq-ssm_document_1.ssm', 'asiaq-ssm_document_2.ssm',
                                        'asiaq-ssm_document_3.ssm', 'asiaq-ssm_document_4.ssm',
                                        'random_file1.txt', 'random_file2.jpg']

        mock_doc_content_1 = '{"random_field_1": "random_value_1"}'
        mock_doc_content_2 = '{"random_field_2": "random_value_2"}'
        mock_file_contents = copy.copy(MOCK_ASIAQ_DOCUMENT_FILE_CONTENTS)
        mock_file_contents[SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_3.ssm'] = mock_doc_content_1
        mock_file_contents[SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_4.ssm'] = mock_doc_content_2
        mock_open.side_effect = create_mock_open(mock_file_contents)

        # Calling the method under test
        self._ssm.update(wait=False)

        # Verify documents are created successfully
        self.assertEqual(_standardize_json_str(mock_doc_content_1),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_3')))
        self.assertEqual(_standardize_json_str(mock_doc_content_2),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_4')))

    @patch('boto3.client', mock_boto3_client)
    @patch('os.listdir')
    @patch('disco_aws_automation.disco_ssm.open')
    def test_update_delete_docs(self, mock_open, mock_os_listdir):
        """Verify that deleting documents in the update() method works"""
        # Setting up test
        mock_os_listdir.return_value = ['asiaq-ssm_document_2.ssm', 'random_file1.txt',
                                        'random_file2.jpg']
        mock_file_contents = copy.copy(MOCK_ASIAQ_DOCUMENT_FILE_CONTENTS)
        mock_file_contents.pop(SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_1.ssm', None)
        mock_open.side_effect = create_mock_open(mock_file_contents)

        # Calling the method under test
        self._ssm.update(wait=False)

        # Verify only document_1 is deleted
        self.assertTrue(self._ssm.get_document_content('asiaq-ssm_document_1') is None)
        self.assertEqual(_standardize_json_str(MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_2']),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_2')))

    @patch('boto3.client', mock_boto3_client)
    @patch('os.listdir')
    @patch('disco_aws_automation.disco_ssm.open')
    def test_update_modify_docs(self, mock_open, mock_os_listdir):
        """Verify that modifying documents in the update() method works"""
        # Setting up test
        mock_os_listdir.return_value = ['asiaq-ssm_document_1.ssm', 'asiaq-ssm_document_2.ssm',
                                        'random_file1.txt', 'random_file2.jpg']

        new_doc_1_content = '{"random_field_1": "random_value_1"}'
        mock_file_contents = copy.copy(MOCK_ASIAQ_DOCUMENT_FILE_CONTENTS)
        mock_file_contents[SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_1.ssm'] = new_doc_1_content
        mock_open.side_effect = create_mock_open(mock_file_contents)

        # Calling the method under test
        self._ssm.update(wait=False)

        # Verify only document_1 is modified
        self.assertEqual(_standardize_json_str(new_doc_1_content),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_1')))
        self.assertEqual(_standardize_json_str(MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_2']),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_2')))

    @patch('boto3.client', mock_boto3_client)
    @patch('os.listdir')
    @patch('disco_aws_automation.disco_ssm.open')
    def test_update_modify_docs_wait(self, mock_open, mock_os_listdir):
        """Verify that modifying documents in the update() method works with wait set to true"""
        # Setting up test
        mock_os_listdir.return_value = ['asiaq-ssm_document_1.ssm', 'asiaq-ssm_document_2.ssm',
                                        'random_file1.txt', 'random_file2.jpg']

        new_doc_1_content = '{"random_field_1": "random_value_1"}'
        mock_file_contents = copy.copy(MOCK_ASIAQ_DOCUMENT_FILE_CONTENTS)
        mock_file_contents[SSM_DOCUMENTS_DIR + '/asiaq-ssm_document_1.ssm'] = new_doc_1_content
        mock_open.side_effect = create_mock_open(mock_file_contents)
        disco_ssm.SSM_WAIT_TIMEOUT = 1
        disco_ssm.SSM_WAIT_SLEEP_INTERVAL = 1

        # Calling the method under test
        self._ssm.update(wait=True)

        # Verify only document_1 is modified
        describe_call = call(Name='asiaq-ssm_document_1')
        # Expecting describe_document() to be called four times: two for delete, two for create
        expected_describe_calls = [describe_call, describe_call, describe_call, describe_call]
        self.assertEqual(expected_describe_calls,
                         self._ssm.conn.describe_document.mock_calls)

        self.assertEqual(_standardize_json_str(new_doc_1_content),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_1')))
        self.assertEqual(_standardize_json_str(MOCK_ASIAQ_DOCUMENT_CONTENTS['asiaq-ssm_document_2']),
                         _standardize_json_str(
                             self._ssm.get_document_content('asiaq-ssm_document_2')))

    @patch('boto3.client', mock_boto3_client)
    def test_get_s3_bucket(self):
        """Verify that we get correct S3 bucket"""

        self.assertEqual(
            MOCK_S3_BUCKET_NAME,
            self._ssm.get_s3_bucket_name()
        )

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_s3_bucket(self):
        """Verify that we get the correct output from an S3 bucket"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            output_s3_bucket_name='foo-bucket'
        )

        command_id = mock_command['CommandId']
        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('stdout', output[0]['stdout'])
            self.assertEqual('stderr', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_s3_bucket_with_no_stdout(self):
        """Verify that we get the correct output from an S3 bucket with no stdout"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            output_s3_bucket_name='foo-bucket',
            stdout=None
        )

        command_id = mock_command['CommandId']
        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('-', output[0]['stdout'])
            self.assertEqual('stderr', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_s3_bucket_with_no_stderr(self):
        """Verify that we get the correct output from an S3 bucket with no stderr"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            output_s3_bucket_name='foo-bucket',
            stderr=None
        )

        command_id = mock_command['CommandId']
        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('stdout', output[0]['stdout'])
            self.assertEqual('-', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_s3_bucket_with_no_output(self):
        """Verify that we get the correct output from an S3 bucket with no output"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            output_s3_bucket_name='foo-bucket',
            stdout=None,
            stderr=None
        )

        command_id = mock_command['CommandId']
        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('-', output[0]['stdout'])
            self.assertEqual('-', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_ssm(self):
        """Verify that we get the correct output from the SSM service"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc'
        )
        command_id = mock_command['CommandId']

        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('stdout', output[0]['stdout'])
            self.assertEqual('stderr', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_ssm_bucket_with_no_stdout(self):
        """Verify that we get the correct output from SSM with no stdout"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            stdout=None
        )
        command_id = mock_command['CommandId']

        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('-', output[0]['stdout'])
            self.assertEqual('stderr', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_ssm_bucket_with_no_stderr(self):
        """Verify that we get the correct output from SSM with no stderr"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            stderr=None
        )
        command_id = mock_command['CommandId']

        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('stdout', output[0]['stdout'])
            self.assertEqual('-', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_get_output_from_ssm_bucket_with_no_output(self):
        """Verify that we get the correct output from SSM with no output"""
        instance_ids = ['i-1', 'i-2']
        mock_command = _create_mock_command(
            instance_ids=instance_ids,
            document_name='foo-doc',
            stderr=None,
            stdout=None
        )
        command_id = mock_command['CommandId']

        command_output = self._ssm.get_ssm_command_output(command_id)

        self.assertEqual(instance_ids, command_output.keys())

        for output in command_output.values():
            self.assertEqual('-', output[0]['stdout'])
            self.assertEqual('-', output[0]['stderr'])

    @patch('boto3.client', mock_boto3_client)
    def test_execute_command(self):
        """Verify that we can execute a command"""
        self._ssm.get_s3_bucket_name = MagicMock(return_value=None)
        instance_ids = ['i-1', 'i-2']
        document_name = "foo-doc"
        comment = "foo-comment"
        parameters = "foo-parameters"

        is_successful = self._ssm.execute(
            instance_ids,
            document_name,
            comment=comment,
            parameters=parameters
        )

        self.assertEqual(True, is_successful)
        self.assertEqual(False, self._ssm.s3.get_object.called)

    @patch('boto3.client', mock_boto3_client)
    def test_execute_command_with_s3(self):
        """Verify that we can execute a command with output in an S3 bucket"""
        instance_ids = ['i-1', 'i-2']
        document_name = "foo-doc"
        comment = "foo-comment"
        parameters = "foo-parameters"

        is_successful = self._ssm.execute(
            instance_ids,
            document_name,
            comment=comment,
            parameters=parameters
        )

        self.assertEqual(True, is_successful)
        self.assertEqual(True, self._ssm.s3.get_object.called)

    @patch('boto3.client', mock_boto3_client)
    def test_execute_command_with_bad_s3(self):
        """Verify that we can execute a command with a bad S3 bucket"""
        self._ssm.get_s3_bucket_name = MagicMock(return_value='asddsa')
        instance_ids = ['i-1', 'i-2']
        document_name = "foo-doc"
        comment = "foo-comment"
        parameters = "foo-parameters"

        is_successful = self._ssm.execute(
            instance_ids,
            document_name,
            comment=comment,
            parameters=parameters
        )

        self.assertEqual(True, is_successful)
        self.assertEqual(True, self._ssm.s3.head_bucket.called)
        self.assertEqual(False, self._ssm.s3.get_object.called)

    @patch('boto3.client', mock_boto3_client)
    def test_execute_command_fails_with_other_status(self):
        """Verify that we fail if the desired status isn't met"""
        self._ssm.get_s3_bucket_name = MagicMock(return_value=None)
        instance_ids = ['i-1', 'i-2']
        document_name = "foo-doc"
        comment = "foo-comment"
        parameters = "foo-parameters"

        is_successful = self._ssm.execute(
            instance_ids,
            document_name,
            comment=comment,
            parameters=parameters,
            desired_status='Failure'
        )

        self.assertEqual(False, is_successful)

    @patch('boto3.client', mock_boto3_client)
    def test_execute_command_with_exception(self):
        """Verify that we fail if there is an exception"""
        self._ssm._send_command = MagicMock(side_effect=ClientError({'Error': {}}, ''))
        self._ssm.get_s3_bucket_name = MagicMock(return_value=None)
        instance_ids = ['i-1', 'i-2']
        document_name = "foo-doc"
        comment = "foo-comment"
        parameters = "foo-parameters"

        is_successful = self._ssm.execute(
            instance_ids,
            document_name,
            comment=comment,
            parameters=parameters
        )

        self.assertEqual(False, is_successful)

    @patch('boto3.client', mock_boto3_client)
    def test_read_env_from_config(self):
        """Verify that we read the env from config if none is provided"""
        config_aws = get_mock_config(MOCK_AWS_CONFIG_DEFINITION)
        self._ssm = DiscoSSM(config_aws=config_aws)

        self.assertEqual(TEST_ENV_NAME, self._ssm.environment_name)
