Go to Qiita Advent Calendar 2024 Top
LoginSignup
6
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

More than 5 years have passed since last update.

AnsibleAdvent Calendar 2019

Day 5
@picapica

AWS関連のansibleの自作モジュールをmotoを使ってunittestしながら作る

Last updated at Posted at 2019-12-04

AWSを使用した自作モジュールを作ることが結構あるのですが、開発の初期段階から直接awsを実行すると時間もかかるしコストもかかります。

なのでできる限りローカルで開発可能なunittestを導入したいと思います。

ここでは自作モジュールを作る時の対象としてなりやすい、aws関連のモジュールのテスト方法を書いていきたいと思います。

機能フロー

  1. インスタンスがrunning状態であれば stoppedの状態にする
  2. stoppedの状態になった時点でインスタンスタイプの変更をする
  3. 起動をし直しrunning状態まで待つ( state: stoppedなら起動はしない)

機能のポイント

インスタンスの起動停止部分

インスタンスの状態を引数として、起動または停止後にステータスが変更されるまで待つ、というロジックにします。

def ensure_state(self,state,wait_time=10,wait_timeout=600):
    passed_time = 0
    if state != self.get_instance_state():
        if state == 'stopped':
            self.changed = True
            retry(lambda: self.client.stop_instances(InstanceIds=[self.instance_id]))
        elif state == 'running':
            self.changed = True
            retry(lambda: self.client.start_instances(InstanceIds=[self.instance_id]))
    while passed_time <=wait_timeout:
        if state == self.get_instance_state():
            return True
        time.sleep(wait_time)
    return False

def get_instance_state(self):
    res = retry(lambda: self.client.describe_instance_status(InstanceIds=[self.instance_id], IncludeAllInstances=True))
    return res["InstanceStatuses"][0]["InstanceState"]["Name"]

インスタンタイプを変更する部分

ここはシンプルにインスタンタイプを変更するだけになります。

def change_instance_type(self):
    self.changed = True
    res = self.client.modify_instance_attribute(InstanceId=self.instance_id,
        InstanceType={'Value':self.new_instance_type})

def check_instance_type(self):
    res = retry(lambda: self.client.describe_instances(InstanceIds=[self.instance_id]))
    return res["Reservations"][0]["Instances"][0]['InstanceType'] == self.new_instance_type

retry処理

上記のawsへの処理に対して retry処理を入れています。
ここでのretry処理は exponential backoffと呼ばれるもので、APIコールが失敗した時に指数的にリトライ時間を延ばしていく、というものになります。
AWSの公式ドキュメントにもリトライ処理を進めるものがあります。
https://docs.aws.amazon.com/ja_jp/general/latest/gr/api-retries.html

よく見るエラーとしては RequestLimitExceededなどですね。こういうエラーに対してのエラーハンドリング処理になります。

def retry(f, retries=5):
    '''
    This is wrapper for exponential backoff
    '''
    exception = None
    for i in range(retries):
        try:
            return f()
        except Exception as e:
            exception = e
            time.sleep(2 ** i)
    raise exception

テストの前準備

ansibleのmoduleをテストする時はいくつかのpatchを事前に差し込む必要があります。
以下のansibleドキュメントにもありますが、

  1. moudleの引数を渡す部分
  2. module自体の終了時に意図的に例外を発生させる

が少なくとも必要になります。
詳しい仕様は公式ドキュメントに書いてあります。
https://docs.ansible.com/ansible/latest/dev_guide/testing_units_modules.html

moudleの引数を渡す部分

ansibleのmoduleの引数は ansible.module_utils.basicで制御をしているため、以下のように引数を渡すための関数を作成します。

def set_module_args(args):
    """prepare arguments so that they will be picked up during module creation"""
    args = json.dumps({'ANSIBLE_MODULE_ARGS': args})
    basic._ANSIBLE_ARGS = to_bytes(args)

module自体の終了時に意図的に例外を発生させる

ansibleのmoduleでは終了時に exit_jsonまたは fail_jsonを実行しています。これらの関数はsys.exit(0)で処理を終了させてしまうため、意図的に例外を発生させることでテストを可能にします。

ポイントは

  • unittestのsetUpでexit_jsonとfail_jsonに対してpatchを当てて、返り値をAnsibleExitJsonかAnsibleFailJsonにする
    です。
class AnsibleExitJson(Exception):
    """Exception class to be raised by module.exit_json and caught by the test case"""
    pass


class AnsibleFailJson(Exception):
    """Exception class to be raised by module.fail_json and caught by the test case"""
    pass


def exit_json(*args, **kwargs):
    """function to patch over exit_json; package return data into an exception"""
    if 'changed' not in kwargs:
        kwargs['changed'] = False
    raise AnsibleExitJson(kwargs)


def fail_json(*args, **kwargs):
    """function to patch over fail_json; package return data into an exception"""
    kwargs['failed'] = True
    raise AnsibleFailJson(kwargs)

class TestMyModule(unittest.TestCase):
    def setUp(self):
        self.mock_module_helper = patch.multiple(basic.AnsibleModule,
                                                 exit_json=exit_json,
                                                 fail_json=fail_json)
        self.mock_module_helper.start()
        self.addCleanup(self.mock_module_helper.stop)

テストケース

ここまで来てやっとテストケースについて考えてみます。

  1. 最終的に指定をしたstateになっているか
  2. インスタンスタイプが変わっているか

この二つをテストしてみましょう。

motoを使ってmockをする

ec2やaws関連のリソースに対してはmotoというlibraryが便利です。
以下のように @mock_ec2というアノテーションをつけるだけで boto3.clientへのAPIアクセスが全てモックされます。

import boto3
from unittest import mock
from moto import mock_ec2

def add_servers(instance_types,server_count=3):
    ami_id = 'ami-1234abcd'
    client = boto3.client('ec2', region_name='ap-northeast-1')
    result = []
    result.append(client.run_instances(
        ImageId=ami_id, 
        InstanceType=instance_types,
        MaxCount=server_count,
        MinCount=server_count
    ))
    return result[0]['Instances'][0]['InstanceId']

    @mock_ec2
    def test_change_instance_types(self):

実際のテストロジック

以下のようなテストロジックになります

  1. moduleの実行
  2. describe_instancesでインスタンス情報を再度取得
  3. 期待した値に変更されているかを確認
class TestMyModule(unittest.TestCase):

    def setUp(self):
        self.mock_module_helper = patch.multiple(basic.AnsibleModule,
                                                 exit_json=exit_json,
                                                 fail_json=fail_json)
        self.mock_module_helper.start()
        self.addCleanup(self.mock_module_helper.stop)

    @mock_ec2
    def test_change_instance_types(self):
        #expected values
        current_instance_type = 'c5.large'
        expected_status = 'running'
        expected_instance_type = 'c5.xlarge'
        region = 'ap-northeast-1'

        #Make ec2 mock
        instance_id = add_servers(current_instance_type)

        #Setup args for ansible
        set_module_args({
            'aws_access_key': 'AWSKEYIOSFODNN7EXAMPLE',
            'aws_secret_key': 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            'instance_id': instance_id,
            'state': expected_status,
            'new_instance_type': expected_instance_type,
            'region': region
        })
        
        #Call actual ansible module
        with self.assertRaises(AnsibleExitJson):
            test_instance_type_change.main()
        
        #Confirm 1. check instance type 2. status is running
        client = boto3.client('ec2', region_name='ap-northeast-1')
        for ins in client.describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances']:
            self.assertEqual(ins['State']['Name'], expected_status)
            self.assertEqual(ins['InstanceType'], expected_instance_type)

実行してみる

問題なくtestが実行されていることが分かります。

root@3560cc651526:/tmp/library# python3 test_module.py
----------------------------------------------------------------------
Ran 1 test in 0.541s

OK

完成したmoudleとテスト

自作モジュール部分

modify_instance_type.py
import time
try:
    import boto3
    import botocore
    HAS_BOTO3_API = True
except ImportError:
    HAS_BOTO3_API = False

from ansible.module_utils.ec2 import aws_common_argument_spec

class ChangeInstanceType(object):
    def __init__(self, module):
        self.changed = False
        self.module = module
        self.aws_access_key = module.params.get('aws_access_key')
        self.aws_secret_key = module.params.get('aws_secret_key')
        self.region = module.params.get('region')
        self.instance_id = module.params.get('instance_id')
        self.new_instance_type = module.params.get('new_instance_type')
        self.state =module.params.get('state')

        if not HAS_BOTO3_API:
                self.module.fail_json(changed=False, msg="Python package boto3 is required")
        try:
            self.client = boto3.client(
                'ec2',
                aws_access_key_id=self.aws_access_key,
                aws_secret_access_key=self.aws_secret_key,
                region_name=self.region
            )
        except botocore.exceptions.ClientError as e:
            self.module.fail_json(changed=self.changed, msg="Cannot initialize connection to ec2: {}".format(e))
    
    def ensure_state(self,state,wait_time=10,wait_timeout=600):
        passed_time = 0
        if state != self.get_instance_state():
            if state == 'stopped':
                self.changed = True
                retry(lambda: self.client.stop_instances(InstanceIds=[self.instance_id]))
            elif state == 'running':
                self.changed = True
                retry(lambda: self.client.start_instances(InstanceIds=[self.instance_id]))
        while passed_time <=wait_timeout:
            if state == self.get_instance_state():
                return True
            time.sleep(wait_time)
        return False
    
    def get_instance_state(self):
        res = retry(lambda: self.client.describe_instance_status(InstanceIds=[self.instance_id], IncludeAllInstances=True))
        return res["InstanceStatuses"][0]["InstanceState"]["Name"]
    
    def change_instance_type(self):
        self.changed = True
        res = self.client.modify_instance_attribute(InstanceId=self.instance_id,
            InstanceType={'Value':self.new_instance_type})
    
    def check_instance_type(self):
        res = retry(lambda: self.client.describe_instances(InstanceIds=[self.instance_id]))
        return res["Reservations"][0]["Instances"][0]['InstanceType'] == self.new_instance_type
    
    def main(self):
        if self.check_instance_type():
            self.changed = False
            return False
        if not self.ensure_state('stopped'):
            self.module.fail_json(changed=self.changed, msg="Module is failed when ensure the state")
        self.change_instance_type()
        if not self.ensure_state(self.state):
            self.module.fail_json(changed=self.changed, msg="Module is failed when ensure the state")

        
def retry(f, retries=5):
    '''
    This is wrapper for exponential backoff
    '''
    exception = None
    for i in range(retries):
        try:
            return f()
        except Exception as e:
            exception = e
            time.sleep(2 ** i)
    raise exception


def main():
    argument_spec = aws_common_argument_spec()
    argument_spec.update(dict(
        aws_access_key=dict(required=True, type='str'),
        aws_secret_key=dict(required=True, type='str', no_log=True),
        region=dict(choices=['us-east-1', 'us-west-2', 'us-west-1', 'eu-west-1', 'eu-central-1', 'ap-southeast-1', 'ap-northeast-1', 'ap-southeast-2', 'ap-northeast-2', 'ap-south-1', 'sa-east-1']),
        instance_id=dict(required=True, type='str'),
        state=dict(type="str", choices=["running", "stopped"], default="running"),
        new_instance_type=dict(required=True, type='str')
    ))
    module = AnsibleModule(argument_spec=argument_spec)
    instance_change_type = ChangeInstanceType(module)
    instance_change_type.main()
    module.exit_json(changed=instance_change_type.changed)

from ansible.module_utils.basic import *
if __name__ == '__main__':
    main()

テスト部分

test_modify_instance_type.py
import json
import sys
import boto3
import botocore
import time
import threading
from pathlib import Path
from unittest import mock
from moto import mock_ec2
from ansible.compat.tests import unittest
from ansible.compat.tests.mock import patch
from ansible.module_utils import basic
from ansible.module_utils._text import to_bytes
import test_instance_type_change

def add_servers(instance_types,server_count=3):
    ami_id = 'ami-1234abcd'
    client = boto3.client('ec2', region_name='ap-northeast-1')
    result = []
    result.append(client.run_instances(
        ImageId=ami_id, 
        InstanceType=instance_types,
        MaxCount=server_count,
        MinCount=server_count
    ))
    return result[0]['Instances'][0]['InstanceId']


def stop_servers(instance_ids):
    client = boto3.client('ec2', region_name='ap-northeast-1')
    client.stop_instances(InstanceIds=instance_ids)


def set_module_args(args):
    """prepare arguments so that they will be picked up during module creation"""
    args = json.dumps({'ANSIBLE_MODULE_ARGS': args})
    basic._ANSIBLE_ARGS = to_bytes(args)


class AnsibleExitJson(Exception):
    """Exception class to be raised by module.exit_json and caught by the test case"""
    pass


class AnsibleFailJson(Exception):
    """Exception class to be raised by module.fail_json and caught by the test case"""
    pass


def exit_json(*args, **kwargs):
    """function to patch over exit_json; package return data into an exception"""
    if 'changed' not in kwargs:
        kwargs['changed'] = False
    raise AnsibleExitJson(kwargs)


def fail_json(*args, **kwargs):
    """function to patch over fail_json; package return data into an exception"""
    kwargs['failed'] = True
    raise AnsibleFailJson(kwargs)

class TestMyModule(unittest.TestCase):

    def setUp(self):
        self.mock_module_helper = patch.multiple(basic.AnsibleModule,
                                                 exit_json=exit_json,
                                                 fail_json=fail_json)
        self.mock_module_helper.start()
        self.addCleanup(self.mock_module_helper.stop)

    @mock_ec2
    def test_change_instance_types(self):
        #expected values
        current_instance_type = 'c5.large'
        expected_status = 'running'
        expected_instance_type = 'c5.xlarge'
        region = 'ap-northeast-1'

        #Make ec2 mock
        instance_id = add_servers(current_instance_type)

        #Setup args for ansible
        set_module_args({
            'aws_access_key': 'AWSKEYIOSFODNN7EXAMPLE',
            'aws_secret_key': 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY',
            'instance_id': instance_id,
            'state': expected_status,
            'new_instance_type': expected_instance_type,
            'region': region
        })
        
        #Call actual ansible module
        with self.assertRaises(AnsibleExitJson):
            test_instance_type_change.main()
        
        #Confirm 1. check instance type 2. status is running
        client = boto3.client('ec2', region_name='ap-northeast-1')
        for ins in client.describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances']:
            self.assertEqual(ins['State']['Name'], expected_status)
            self.assertEqual(ins['InstanceType'], expected_instance_type)
 
if __name__ == "__main__":
    unittest.main()

おわりに

ansibleのテストは標準のunittestで実行できないものが結構あるので厄介ですね。
ansibleの公式のrepositoryでも各moduleに対してのテストがコミットされているので、それを参考に自作モジュールを作っていくのも良さそうです。

6
3
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
6
3

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?