actix-web でリトライ時の tracing のログをテストしてみた

動機

あるアプリケーション内に外部APIを実行する処理がある。
そのAPIには短時間でのレートリミットが存在するため、その場合は即座にエラーとするのではなく何度かリトライさせる設計にしている。

この関数はあくまで外部APIの実行を責務とし、リトライが発生したことはログを出力することで気づけるようにしている。
リトライした回数など重要な情報がログに含まれているため、リトライ処理の自動テストで出力されているログの内容についても確認したい。

やりたいこと

• リトライの有／無、リトライを上限まで実施した場合などの網羅的なテストコードを用意したい
• tracing のログ内容もテストしたい

コード

外部APIを実行する処理のサンプル

use reqwest::StatusCode;

async fn execute_api(endpoint_url: &str) -> Result<serde_json::Value, String> {
    let client = reqwest::Client::new();
    let try_count = 3;

    for _ in 0..try_count {
        let res = client.post(endpoint_url).send().await.unwrap();

        // 200 OK
        if res.status().is_success() {
            let body_bytes = res.bytes().await.unwrap();
            return Ok(serde_json::from_slice(&body_bytes).unwrap());
        }

        // 429 TOO_MANY_REQUESTS => リトライ
        if res.status() == reqwest::StatusCode::TOO_MANY_REQUESTS {
            tracing::warn!("レートリミットに抵触しました。");
            continue;
        }
    }

    tracing::error!("リトライ回数を超過しました。");
    Err("リトライ上限".to_string())
}

テストコードのサンプル

#[cfg(test)]
mod tests {
    use super::*;
    use mockito::ServerGuard;
    use std::{
        io::Write,
        sync::{Arc, Mutex},
    };
    use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer};
    use tracing_subscriber::{fmt::MakeWriter, layer::SubscriberExt, EnvFilter};

    struct StringWriter {
        buf: Arc<Mutex<Vec<String>>>,
    }

    impl StringWriter {
        fn new(buf: &Arc<Mutex<Vec<String>>>) -> Self {
            Self {
                buf: Arc::clone(&buf),
            }
        }
    }

    impl Write for StringWriter {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            if let Ok(mut buffer) = self.buf.lock() {
                buffer.push(String::from_utf8_lossy(buf).to_string());
            }
            Ok(buf.len())
        }

        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    impl<'a> MakeWriter<'a> for StringWriter {
        type Writer = Self;

        fn make_writer(&'a self) -> Self::Writer {
            StringWriter {
                buf: Arc::clone(&self.buf),
            }
        }
    }

    fn setup_subscriber(name: impl Into<String>) -> (impl tracing::Subscriber, Arc<Mutex<Vec<String>>>) {
        let buffer = Arc::new(Mutex::new(vec![]));

        let formatting_layer = BunyanFormattingLayer::new(name.into(), StringWriter::new(&buffer));
        let subscriber = tracing_subscriber::registry()
            .with(JsonStorageLayer)
            .with(formatting_layer)
            // レベル50以上のログだけキャプチャ
            .with(EnvFilter::new("warn"));

        (subscriber, buffer)
    }

    // 1回リトライして2回目で成功するテスト
    #[tokio::test]
    async fn test_request_succeed_on_second_try() {
        let (subscriber, buffer) = setup_subscriber("test_request_succeed_on_second_try");
        let _guard = tracing::subscriber::set_default(subscriber);

        let mock_server = {
            let mut mock_server = mockito::Server::new_async().await;

            // 1回目のリクエストは 429 を返す
            let _mock = mock_server
                .mock("POST", "/")
                .with_header("Content-Type", "application/json")
                .with_body(r#"{"message": "Too Many Requests"}"#)
                .with_status(429)
                .expect(1)
                .create();
            // 2回目のリクエストは 200 を返す
            let _mock = mock_server
                .mock("POST", "/")
                .with_header("Content-Type", "application/json")
                .with_body(r#"{"message": "ok"}"#)
                .with_status(200)
                .expect(1)
                .create();

            mock_server
        };

        let endpoint_url = format!("{}/", mock_server.url());
        let res = execute_api(&endpoint_url).await.unwrap();
        assert!(res.is_ok(), "{:?}", res.unwrap_err());

        let logs = {
            (*buffer.lock().unwrap())
                .iter()
                .map(|log| serde_json::from_str(log).unwrap())
                .collect()
        };

        // ログはリトライ1回分のみ
        assert_eq!(logs.len(), 1);

        // リトライログのメッセージが正しいか
        assert_eq!(&logs[0]["msg"], "レートリミットに抵触しました。");
    }
}

• ログ内容をテストできるよう独自の Writer を用意し、ログの出力先としてセットした
  • tracing::subscriber::set_default の実行はテスト関数内で実行しないとキャプチャできない点に注意
• mockito を使用してリクエストごとに任意のレスポンスを返すようにすることで、リトライの実施状況別のテスト関数を用意してログのテストをすることができる