from unittest.mock import MagicMock

import boto
import boto3

from processing_accounting.util import emr


def test_launch_emr_cluster(monkeypatch):
    mock_emr_client = MagicMock()

    monkeypatch.setattr(
        mock_emr_client, 'run_job_flow',
        MagicMock(return_value=dict(JobFlowId='test_cluster_id')))

    mock_activity = MagicMock()
    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        boto3, 'client', MagicMock(return_value=mock_emr_client))

    resp = emr.launch_emr_cluster(
        mock_activity, None, 'path_to_logs', 'job_name', 'tag_name',
        'ec2_instance_name', False)

    instances = {
        'Ec2KeyName': 'orchard_admin',
        'KeepJobFlowAliveWhenNoSteps': False,
        'Ec2SubnetId': 'subnet-cb4dbae0',
        'InstanceGroups': emr.INSTANCE_GROUPS
    }

    mock_emr_client.run_job_flow.assert_any_call(
        Name='JOB_NAME',
        ReleaseLabel='emr-5.30.2',
        LogUri='path_to_logs',
        VisibleToAllUsers=True,
        ServiceRole='EMR_DefaultRole',
        JobFlowRole='EMR_EC2_DefaultRole',
        Instances=instances,
        Applications=[dict(Name='hive')])

    tags = []
    tags.append(dict(Key='Name', Value='ec2_instance_name'))
    tags.append(dict(Key='tag_name'))
    mock_emr_client.add_tags.assert_any_call(
        ResourceId='test_cluster_id', Tags=tags)

    assert resp['emr.cluster_id'] == 'test_cluster_id'


def test_add_streaming_job_step(monkeypatch):
    mock_connector = MagicMock()
    mock_response = MagicMock()
    mock_item = MagicMock()
    mock_item.value = 'as78dfasdf9'
    mock_response.stepids = [mock_item]

    monkeypatch.setattr(
        mock_connector, 'run_jobflow',
        MagicMock(return_value='test_cluster_id'))

    monkeypatch.setattr(
        mock_connector, 'add_jobflow_steps',
        MagicMock(return_value=mock_response))

    mock_activity = MagicMock()
    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        boto.emr, 'connect_to_region', MagicMock(return_value=mock_connector))

    resp = emr.add_streaming_job_step(
        mock_activity, 'test_cluster_id', 'step_name', [], 'mapper',
        'reducer', 'input',
        'output', [])

    assert resp['emr.step_id'] == 'as78dfasdf9'


def test_add_s3_distcp_step(monkeypatch):
    mock_connector = MagicMock()
    mock_response = MagicMock()
    mock_item = MagicMock()
    mock_item.value = '868868686886'
    mock_response.stepids = [mock_item]

    monkeypatch.setattr(
        boto.emr, 'connect_to_region', MagicMock(return_value=mock_connector))

    monkeypatch.setattr(
        mock_connector, 'add_jobflow_steps',
        MagicMock(return_value=mock_response))

    mock_activity = MagicMock()
    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    resp = emr.add_s3_distcp_step(
        mock_activity, 'emr_cluster_id', 'source_path',
        'destination_path')

    assert resp['emr.step_id'] == '868868686886'


def test_add_hive_job_step(monkeypatch):
    mock_connector = MagicMock()
    mock_response = MagicMock()
    mock_item = MagicMock()
    mock_item.value = '868868686886'
    mock_response.stepids = [mock_item]

    mock_activity = MagicMock()
    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        mock_activity, 'info', MagicMock(return_value=None))

    monkeypatch.setattr(
        boto.emr, 'connect_to_region', MagicMock(return_value=mock_connector))

    monkeypatch.setattr(
        mock_connector, 'add_jobflow_steps',
        MagicMock(return_value=mock_response))

    resp = emr.add_hive_job_step(
        mock_activity, 'emr_cluster_id', 'source_s3_path', 'step_name',
        'destination_s3_path', 'hql_s3_path')

    assert resp['emr.step_id'] == '868868686886'