X-RayのSAMとLambda内への組み込みとjestでのモック化

2017年頃からサーバレスに浸っている、TISの小西啓介です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/40049/a3cf8470-1e9b-e715-9505-7a6f5a123e3c.png)

# 概要

SAM(Serverless Application Model)を使ったnode.jsのLambdaプログラム等に対して、
パフォーマンス計測等を行いX-Ray機能の組み込みについて説明します。

また、X-Ray機能を組み込むと、aws-sdk-mockを使用したユニットテストでは
通らなくなるので、aws-xray-sdkのモック化についても説明します。

# SAMへのX-Rayの追加

API Gateway - Lambda - DynamoDBのような構成だとして、
SAM(Serverless Application Model)へのX-Rayの組み込みは、
以下のように、SAMのテンプレーレートのGlobalセクションの
`Function` に対して、`Tracing: Active`と
`Api` に対して、`TracingEnabled: True`を追加するだけです。

もちろん、Globalではなく、AWS::Serverless::Functionや
AWS::Serverless::ApiのPropertiesに追加しても良いです。

```yaml:変更前SAM-Template

Globals:
  Function:
    Timeout: 15
  Api:
    # https://github.com/awslabs/serverless-application-model/issues/191
    OpenApiVersion: 3.0.0
```

```yaml:変更後SAM-Template
Globals:
  Function:
    Timeout: 15
    Tracing: Active # 追加
  Api:
    # https://github.com/awslabs/serverless-application-model/issues/191
    OpenApiVersion: 3.0.0
    TracingEnabled: True # 追加
```

# Lambda内のAWSサービス呼出へのX-Rayの組み込み

API GatewayからLambdaへの呼出のみであれば、上記のSAMへの組み込みのみで十分なのですが、
Lambda内で、AWS-SDKを使ってAWSのサービスを呼び出している場合、その呼出のトレースを行う場合、コードの変更が必要です。
ただ、コードの変更と言っても、 `npm i -S aws-xray-sdk-core` で`aws-xray-sdk` を組み込んで以下のコードの変更を行うだけです。

```js:変更前js
const AWS = require('aws-sdk');
```

```js:変更後js
const AWSXRay = require('aws-xray-sdk-core');
const AWS = AWSXRay.captureAWS(require('aws-sdk'));
```
なお、`aws-xray-sdk`には、 以下の４つのモジュールが含まれます。

- aws-xray-sdk-core
- aws-xray-sdk-express
- aws-xray-sdk-mysql
- aws-xray-sdk-postgres

Lambdaの場合、mysqlやpostgresに接続しないのであれば、`aws-xray-sdk-core`で十分だと思います。
Lambda内で、[AWS Serverless Express](https://github.com/awslabs/aws-serverless-express)を使用している場合、`aws-xray-sdk-express` が気になりますが、これはEC2用の為、使用できません（Lambdaでは、呼出時にX-Rayのセグメントが既に作られているので、新たにセグメントを設定することは出来ないため）

# ユニットテスト(jest)でのモック化

上記のようなAWS-SDKを使用しているプログラムのユニットテストで便利な [AWS-SDK-Mock](https://github.com/dwyl/aws-sdk-mock)があるのですが、AWS-SDKをラップしている関係で、そのままのコードでは動かなくなってしまします。

AWSXRayが使用する部分AWS-SDKの処理をAWS-SDK-Mockで書けば良いのもしれないのですが、情報が無くよくわかりませんでした。
(`AWS_XRAY_CONTEXT_MISSING` などの環境変数を設定してもだめでした)

また、Jest用に[aws-xray-Mock](https://www.npmjs.com/package/aws-xray-mock)というのが有るのですが、上記のような呼び出し方(captureAWS)には対応しておらず、個別の(AWSサービス毎の)クライアントにしか対応していないようです。プルリクも出ていますが、５月から、放置されてしまっています。
https://github.com/MechanicalRock/aws-xray-mock/pull/4

そこで、上記のコードを参考に、captureAWSのみに対応する場合、
以下の追加で、aws-sdk-mockを使ったテストが全て通るようになりました。

```js
beforeAll(async () => {
    jest.mock('aws-xray-sdk-core', () => {
        return { captureAWS: aws => aws };
    });
});
```

サブセグメントの追加やAWSクライアント(サービス)単位の設定等を行う場合、
以下のように、利用するメソッドを追加しましょう。
(出典：[aws-xray-mock](https://github.com/MechanicalRock/aws-xray-mock) 及び [aws-xray-mockのプルリク](https://github.com/MechanicalRock/aws-xray-mock/pull/4))

```js
let xRaySegment;
class MockXraySegment {
    constructor(name) {
        this.name = name;
        xRaySegment = this;
    }
    addNewSubsegment(name) {
        return new MockXraySegment(name);
    }
    addError(error) { }
    close() { }
}
beforeAll(async () => {
    // process.env.AWS_XRAY_CONTEXT_MISSING = 'LOG_ERROR';
    jest.mock('aws-xray-sdk-core', () => {
        return {
            captureAWS: aws => aws,
            captureAWSClient: client => client,
            captureHTTPs: http => http,
            captureHTTPsGlobal: httpGlobal => httpGlobal,
            getSegment: () => {
                if (!xRaySegment) {
                    xRaySegment = new MockXraySegment('Parent Segment');
                    return xRaySegment;
                }
                return xRaySegment;
            }
        };
    });
});
beforeEach(async () => {
    xRaySegment = null;
});

```

やはり、 `aws-xray-mock` のプルリクマージして欲しいですね。
(`aws-xray-sdk`だけでなくて、`aws-xray-sdk-core` もモック化してほしいですが。。)

おわり

# 参考

- [AWS X-Ray を使っていると aws-mock から\`Service.prototype.customizeRequests is not a function\`と怒られる件 - Qiita](https://qiita.com/snaka/items/7e0e4148ac6b994c2c36)
- [AWS SDKを使ってもAWS環境にアクセスせずにユニットテストを回す方法 - Qiita](https://qiita.com/horike37/items/41f9586d315c91d49305)