import unittest
import json
import boto3
import time

# time delay before removing the alarms for visual inspection via console
wait_seconds_before_termination = 10

function = __import__('lambda_function')
cw = boto3.client('cloudwatch')


class TestFunction(unittest.TestCase):

    def test_function_running_state(self):
        event = load_event_json('event_running.json')
        result = function.lambda_handler(event, {})
        self.assertEqual(result['statusCode'], 200)
        r = cw.describe_alarms(AlarmNames=[
            event['detail']['instance-id'] + '_CPUUtilization',
            event['detail']['instance-id'] + '_disk_used_percent',
            event['detail']['instance-id'] + '_mem_used_percent'
        ])
        self.assertEqual(len(r['MetricAlarms']), 3)

    def test_function_terminated_state(self):
        time.sleep(wait_seconds_before_termination)
        event = load_event_json('event_terminated.json')
        result = function.lambda_handler(event, {})
        self.assertEqual(result['statusCode'], 200)
        r = cw.describe_alarms(AlarmNames=[
            event['detail']['instance-id'] + '_CPUUtilization',
            event['detail']['instance-id'] + '_disk_used_percent',
            event['detail']['instance-id'] + '_mem_used_percent'
        ])
        self.assertEqual(len(r['MetricAlarms']), 0)


def load_event_json(filepath):
    with open(filepath, 'r') as fd:
        return json.load(fd)


if __name__ == '__main__':
    unittest.main()
