summaryrefslogtreecommitdiffstats
path: root/crates/tests/mqtt_tests/src/test_mqtt_client.rs
blob: b4fe0fa40e2e6b0427111b3f15e8746477e67d88 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
use crate::with_timeout::WithTimeout;
use rumqttc::{AsyncClient, Event, EventLoop, MqttOptions, Packet, QoS};
use std::time::Duration;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};

/// Returns the stream of messages received on a specific topic.
///
/// To ease testing, the errors are returned as messages.
pub async fn messages_published_on(mqtt_port: u16, topic: &str) -> UnboundedReceiver<String> {
    let (sender, recv) = tokio::sync::mpsc::unbounded_channel();

    // One can have a connection error if this is called just after the broker starts
    // So try to subscribe again after a first error
    let mut con = TestCon::new(mqtt_port);
    let mut retry = 1;
    loop {
        match con.subscribe(topic, QoS::AtLeastOnce).await {
            Ok(()) => break,
            Err(_) if retry > 0 => {
                tokio::time::sleep(Duration::from_secs(1)).await;
                retry -= 1;
                continue;
            }
            Err(err) => {
                let msg = format!("Error: {:?}", err).to_string();
                let _ = sender.send(msg);
                return recv;
            }
        }
    }

    tokio::spawn(async move {
        con.forward_received_messages(sender).await;
    });

    recv
}

/// Check that a list of messages has been received in the given order
pub async fn assert_received<T>(
    messages: &mut UnboundedReceiver<String>,
    timeout: Duration,
    expected: T,
) where
    T: IntoIterator,
    T::Item: ToString,
{
    for expected_msg in expected.into_iter() {
        let actual_msg = messages.recv().with_timeout(timeout).await;
        assert_eq!(actual_msg, Ok(Some(expected_msg.to_string())));
    }
}

/// Publish a message
///
/// Return only when the message has been acknowledged.
pub async fn publish(
    mqtt_port: u16,
    topic: &str,
    payload: &str,
    qos: QoS,
    retain: bool,
) -> Result<(), anyhow::Error> {
    let mut con = TestCon::new(mqtt_port);

    con.publish(topic, qos, retain, payload).await
}

/// Publish the `pub_message` on the `pub_topic` only when ready to receive a message on `sub_topic`.
///
/// 1. Subscribe to the `sub_topic`,
/// 2. Wait for the acknowledgment of the subscription
/// 3  Publish the `pub_message` on the `pub_topic`,
/// 4. Return the first received message
/// 5. or give up after `timeout_sec` secondes.
pub async fn wait_for_response_on_publish(
    mqtt_port: u16,
    pub_topic: &str,
    pub_message: &str,
    sub_topic: &str,
    timeout: Duration,
) -> Option<String> {
    let mut con = TestCon::new(mqtt_port);

    con.subscribe(sub_topic, QoS::AtLeastOnce).await.ok()?;
    con.publish(pub_topic, QoS::AtLeastOnce, false, pub_message)
        .await
        .ok()?;
    match tokio::time::timeout(timeout, con.next_message()).await {
        // One collapse both timeout and error to None
        Err(_) | Ok(Err(_)) => None,
        Ok(Ok(x)) => Some(x),
    }
}

pub async fn map_messages_loop<F>(mqtt_port: u16, func: F)
where
    F: Send + Sync + Fn((String, String)) -> Vec<(String, String)>,
{
    let mut con = TestCon::new(mqtt_port);
    con.subscribe("#", QoS::AtLeastOnce)
        .await
        .expect("Fail to subscribe on #");

    loop {
        if let Ok(message) = con.next_topic_payload().await {
            dbg!(&message);
            for (topic, response) in func(message).iter() {
                let _ = con.publish(topic, QoS::AtLeastOnce, false, response).await;
            }
        }
    }
}

pub struct TestCon {
    client: AsyncClient,
    eventloop: EventLoop,
}

impl TestCon {
    pub fn new(mqtt_port: u16) -> TestCon {
        let id: String = std::iter::repeat_with(fastrand::alphanumeric)
            .take(10)
            .collect();
        let mut options = MqttOptions::new(id, "localhost", mqtt_port);
        options.set_clean_session(true);

        let (client, eventloop) = AsyncClient::new(options, 10);
        TestCon { client, eventloop }
    }

    pub async fn subscribe(&mut self, topic: &str, qos: QoS) -> Result<(), anyhow::Error> {
        self.client.subscribe(topic, qos).await?;

        loop {
            match self.eventloop.poll().await {
                Ok(Event::Incoming(Packet::SubAck(_))) => {
                    return Ok(());
                }
                Err(err) => {
                    return Err(err)?;
                }
                _ => {}
            }
        }
    }

    pub async fn publish(
        &mut self,
        topic: &str,
        qos: QoS,
        retain: bool,
        payload: &str,
    ) -> Result<(), anyhow::Error> {
        self.client.publish(topic, qos, retain, payload).await?;

        loop {
            match self.eventloop.poll().await {
                Ok(Event::Incoming(Packet::PubAck(_))) => {
                    return Ok(());
                }
                Err(err) => {
                    return Err(err)?;
                }
                _ => {}
            }
        }
    }

    pub async fn forward_received_messages(&mut self, sender: UnboundedSender<String>) {
        loop {
            match self.eventloop.poll().await {
                Ok(Event::Incoming(Packet::Publish(response))) => {
                    let msg = std::str::from_utf8(&response.payload)
                        .unwrap_or("Error: non-utf8-payload")
                        .to_string();
                    if let Err(_) = sender.send(msg) {
                        break;
                    }
                }
                Err(err) => {
                    let msg = format!("Error: {:?}", err).to_string();
                    let _ = sender.send(msg);
                    break;
                }
                _ => {}
            }
        }
        let _ = self.client.disconnect().await;
    }

    pub async fn next_message(&mut self) -> Result<String, anyhow::Error> {
        loop {
            match self.eventloop.poll().await {
                Ok(Event::Incoming(Packet::Publish(packet))) => {
                    let msg = std::str::from_utf8(&packet.payload)
                        .unwrap_or("Error: non-utf8-payload")
                        .to_string();
                    return Ok(msg);
                }
                Err(err) => {
                    return Err(err)?;
                }
                _ => {}
            }
        }
    }

    pub async fn next_topic_payload(&mut self) -> Result<(String, String), anyhow::Error> {
        loop {
            match self.eventloop.poll().await {
                Ok(Event::Incoming(Packet::Publish(packet))) => {
                    let topic = packet.topic.clone();
                    let msg = std::str::from_utf8(&packet.payload)
                        .unwrap_or("Error: non-utf8-payload")
                        .to_string();
                    return Ok((topic, msg));
                }
                Err(err) => {
                    return Err(err)?;
                }
                _ => {}
            }
        }
    }
}