aws_smithy_eventstream/
test_util.rs

1/*
2 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6use crate::frame::{read_message_from, DecodedFrame, MessageFrameDecoder};
7use aws_smithy_types::event_stream::{HeaderValue, Message};
8use std::collections::{BTreeMap, BTreeSet};
9use std::error::Error as StdError;
10
11/// Validate that the bodies match, which includes headers and messages
12///
13/// When `full_stream` is true, it also verifies the length of frames
14pub fn validate_body(
15    expected_body: &[u8],
16    actual_body: &[u8],
17    full_stream: bool,
18) -> Result<(), Box<dyn StdError>> {
19    let expected_frames = decode_frames(expected_body);
20    let actual_frames = decode_frames(actual_body);
21
22    if full_stream {
23        assert_eq!(
24            expected_frames.len(),
25            actual_frames.len(),
26            "Frame count didn't match.\n\
27        Expected: {expected_frames:?}\n\
28        Actual:   {actual_frames:?}",
29        );
30    }
31
32    for ((expected_wrapper, expected_message), (actual_wrapper, actual_message)) in
33        expected_frames.into_iter().zip(actual_frames.into_iter())
34    {
35        assert_eq!(
36            header_names(&expected_wrapper),
37            header_names(&actual_wrapper)
38        );
39        if let Some(expected_message) = expected_message {
40            let actual_message = actual_message.unwrap();
41            assert_eq!(header_map(&expected_message), header_map(&actual_message));
42            assert_eq!(expected_message.payload(), actual_message.payload());
43        }
44    }
45    Ok(())
46}
47
48// Returned tuples are (SignedWrapperMessage, WrappedMessage).
49// Some signed messages don't have payloads, so in those cases, the wrapped message will be None.
50fn decode_frames(mut body: &[u8]) -> Vec<(Message, Option<Message>)> {
51    let mut result = Vec::new();
52    let mut decoder = MessageFrameDecoder::new();
53    while let DecodedFrame::Complete(msg) = decoder.decode_frame(&mut body).unwrap() {
54        let inner_msg = if msg.payload().is_empty() {
55            None
56        } else {
57            Some(read_message_from(msg.payload().as_ref()).unwrap())
58        };
59        result.push((msg, inner_msg));
60    }
61    result
62}
63
64fn header_names(msg: &Message) -> BTreeSet<String> {
65    msg.headers()
66        .iter()
67        .map(|h| h.name().as_str().into())
68        .collect()
69}
70fn header_map(msg: &Message) -> BTreeMap<String, &HeaderValue> {
71    msg.headers()
72        .iter()
73        .map(|h| (h.name().as_str().to_string(), h.value()))
74        .collect()
75}
aws_smithy_eventstream/test_util.rs

aws_smithy_eventstream/
test_util.rs
