Coverage Report

Created: 2026-07-21 15:28

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/build/source/nativelink-redis-tester/src/fake_redis.rs
Line
Count
Source
1
// Copyright 2026 The NativeLink Authors. All rights reserved.
2
//
3
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//    See LICENSE file for details
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
15
use core::fmt::Write;
16
use core::hash::BuildHasher;
17
use std::collections::HashMap;
18
19
use nativelink_util::background_spawn;
20
use redis::Value;
21
use redis_test::IntoRedisValue;
22
use tokio::io::{AsyncReadExt, AsyncWriteExt};
23
use tokio::net::TcpListener;
24
use tokio::sync::oneshot::{self, Sender};
25
use tracing::{error, info, warn};
26
27
109
fn cmd_as_string(cmd: &redis::Cmd) -> String {
28
109
    let raw = cmd.get_packed_command();
29
109
    String::from_utf8(raw).unwrap()
30
109
}
31
32
992
pub(crate) fn arg_as_string(output: &mut String, arg: Value) {
33
992
    match arg {
34
117
        Value::SimpleString(s) => {
35
117
            write!(output, "+{s}\r\n").unwrap();
36
117
        }
37
147
        Value::Okay => {
38
147
            write!(output, "+OK\r\n").unwrap();
39
147
        }
40
212
        Value::BulkString(s) => {
41
212
            write!(
42
212
                output,
43
212
                "${}\r\n{}\r\n",
44
212
                s.len(),
45
212
                str::from_utf8(&s).unwrap()
46
212
            )
47
212
            .unwrap();
48
212
        }
49
206
        Value::Int(v) => {
50
206
            write!(output, ":{v}\r\n").unwrap();
51
206
        }
52
192
        Value::Array(values) => {
53
192
            write!(output, "*{}\r\n", values.len()).unwrap();
54
376
            for value in 
values192
{
55
376
                arg_as_string(output, value);
56
376
            }
57
        }
58
65
        Value::Map(values) => {
59
65
            write!(output, "%{}\r\n", values.len()).unwrap();
60
86
            for (key, value) in 
values65
{
61
86
                arg_as_string(output, key);
62
86
                arg_as_string(output, value);
63
86
            }
64
        }
65
52
        Value::Nil => {
66
52
            write!(output, "_\r\n").unwrap();
67
52
        }
68
1
        Value::Boolean(value) => {
69
1
            if value {
70
1
                write!(output, "#t\r\n")
71
            } else {
72
0
                write!(output, "#f\r\n")
73
            }
74
1
            .unwrap();
75
        }
76
        _ => {
77
0
            panic!("No support for {arg:?}")
78
        }
79
    }
80
992
}
81
82
89
fn args_as_string(args: Vec<Value>) -> String {
83
89
    let mut output = String::new();
84
215
    for arg in 
args89
{
85
215
        arg_as_string(&mut output, arg);
86
215
    }
87
89
    output
88
89
}
89
90
49
pub fn add_to_response<B: BuildHasher>(
91
49
    response: &mut HashMap<String, String, B>,
92
49
    cmd: &redis::Cmd,
93
49
    args: Vec<Value>,
94
49
) {
95
49
    add_to_response_raw(response, cmd, args_as_string(args));
96
49
}
97
98
49
pub fn add_to_response_raw<B: BuildHasher>(
99
49
    response: &mut HashMap<String, String, B>,
100
49
    cmd: &redis::Cmd,
101
49
    args: String,
102
49
) {
103
49
    response.insert(cmd_as_string(cmd), args);
104
49
}
105
106
20
fn setinfo(responses: &mut HashMap<String, String>) {
107
    // We do raw inserts of command here, because the library sends 3/4 commands in one go
108
    // They always start with HELLO, then optionally SELECT, so we use this to differentiate
109
20
    let hello = cmd_as_string(redis::cmd("HELLO").arg("3"));
110
20
    let setinfo = cmd_as_string(
111
20
        redis::cmd("CLIENT")
112
20
            .arg("SETINFO")
113
20
            .arg("LIB-NAME")
114
20
            .arg("redis-rs"),
115
    );
116
20
    responses.insert(
117
20
        [hello.clone(), setinfo.clone()].join(""),
118
20
        args_as_string(vec![
119
20
            Value::Map(vec![(
120
20
                Value::SimpleString("server".into()),
121
20
                Value::SimpleString("redis".into()),
122
20
            )]),
123
20
            Value::Okay,
124
20
            Value::Okay,
125
        ]),
126
    );
127
20
    responses.insert(
128
20
        [hello, cmd_as_string(redis::cmd("SELECT").arg(3)), setinfo].join(""),
129
20
        args_as_string(vec![
130
20
            Value::Map(vec![(
131
20
                Value::SimpleString("server".into()),
132
20
                Value::SimpleString("redis".into()),
133
20
            )]),
134
20
            Value::Okay,
135
20
            Value::Okay,
136
20
            Value::Okay,
137
        ]),
138
    );
139
20
}
140
141
13
pub fn add_lua_script<B: BuildHasher>(
142
13
    responses: &mut HashMap<String, String, B>,
143
13
    lua_script: &str,
144
13
    hash: &str,
145
13
) {
146
13
    add_to_response(
147
13
        responses,
148
13
        redis::cmd("SCRIPT").arg("LOAD").arg(lua_script),
149
13
        vec![hash.into_redis_value()],
150
    );
151
13
}
152
153
13
pub fn fake_redis_stream() -> HashMap<String, String> {
154
13
    let mut responses = HashMap::new();
155
13
    setinfo(&mut responses);
156
    // Does setinfo as well, so need to respond to all 3
157
13
    add_to_response(
158
13
        &mut responses,
159
13
        redis::cmd("SELECT").arg("3"),
160
13
        vec![Value::Okay, Value::Okay, Value::Okay],
161
    );
162
13
    responses
163
13
}
164
165
3
pub fn fake_redis_sentinel_master_stream() -> HashMap<String, String> {
166
3
    let mut response = fake_redis_stream();
167
3
    add_to_response(
168
3
        &mut response,
169
3
        &redis::cmd("ROLE"),
170
3
        vec![Value::Array(vec![
171
3
            "master".into_redis_value(),
172
3
            0.into_redis_value(),
173
3
            Value::Array(vec![]),
174
3
        ])],
175
    );
176
3
    response
177
3
}
178
179
7
pub fn fake_redis_sentinel_stream(master_name: &str, redis_port: u16) -> HashMap<String, String> {
180
7
    let mut response = HashMap::new();
181
7
    setinfo(&mut response);
182
183
    // Not a full "sentinel masters" response, but enough for redis-rs
184
7
    let resp: Vec<(Value, Value)> = vec![
185
7
        ("name".into_redis_value(), master_name.into_redis_value()),
186
7
        ("ip".into_redis_value(), "127.0.0.1".into_redis_value()),
187
7
        (
188
7
            "port".into_redis_value(),
189
7
            i64::from(redis_port).into_redis_value(),
190
7
        ),
191
7
        ("flags".into_redis_value(), "master".into_redis_value()),
192
    ];
193
194
7
    add_to_response(
195
7
        &mut response,
196
7
        redis::cmd("SENTINEL").arg("MASTERS"),
197
7
        vec![Value::Array(vec![Value::Map(resp)])],
198
    );
199
7
    response
200
7
}
201
202
28
pub(crate) async fn fake_redis_internal<H>(
203
28
    listener: TcpListener,
204
28
    listener_ready_tx: Sender<()>,
205
28
    handlers: Vec<H>,
206
28
) where
207
28
    H: Fn(&[u8]) -> String + Send + Clone + 'static + Sync,
208
28
{
209
28
    let mut handler_iter = handlers.iter().cloned().cycle();
210
28
    info!(
211
        "Waiting for connection on {}",
212
28
        listener.local_addr().unwrap()
213
    );
214
28
    listener_ready_tx
215
28
        .send(())
216
28
        .expect("Expected successful send");
217
    loop {
218
70
        let Ok((
mut stream42
, _)) = listener.accept().await else {
219
0
            error!("accept error");
220
0
            panic!("error");
221
        };
222
42
        info!("Accepted new connection");
223
42
        let local_handler = handler_iter.next().unwrap();
224
42
        background_spawn!("thread", async move {
225
            loop {
226
103k
                let mut buf = vec![0; 8192];
227
103k
                let 
res103k
= stream.read(&mut buf).await.
unwrap103k
();
228
103k
                if res != 0 {
229
526
                    let output = local_handler(&buf[..res]);
230
526
                    if !output.is_empty() {
231
262
                        stream.write_all(output.as_bytes()).await.unwrap();
232
264
                    }
233
103k
                }
234
            }
235
        });
236
    }
237
}
238
239
21
async fn fake_redis<B>(
240
21
    listener: TcpListener,
241
21
    listener_ready_tx: Sender<()>,
242
21
    all_responses: Vec<HashMap<String, String, B>>,
243
21
) where
244
21
    B: BuildHasher + Clone + Send + 'static + Sync,
245
21
{
246
21
    let funcs = all_responses
247
21
        .iter()
248
21
        .map(|responses| {
249
21
            info!("Responses are: {:?}", responses);
250
21
            let values = responses.clone();
251
329
            move |buf: &[u8]| -> String {
252
329
                let str_buf = String::from_utf8_lossy(buf).into_owned();
253
1.24k
                for (key, value) in 
&values329
{
254
1.24k
                    if str_buf.starts_with(key) {
255
65
                        info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n"));
256
65
                        return value.clone();
257
1.17k
                    }
258
                }
259
264
                warn!(
260
                    "Unknown command: {}",
261
264
                    str_buf.chars().take(1000).collect::<String>()
262
                );
263
264
                String::new()
264
329
            }
265
21
        })
266
21
        .collect();
267
21
    fake_redis_internal(listener, listener_ready_tx, funcs).await;
268
0
}
269
270
21
async fn make_fake_redis_with_multiple_responses<B: BuildHasher + Clone + Send + 'static + Sync>(
271
21
    responses: Vec<HashMap<String, String, B>>,
272
21
) -> u16 {
273
21
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
274
21
    let port = listener.local_addr().unwrap().port();
275
276
21
    let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>();
277
278
21
    background_spawn!("listener", async move {
279
21
        fake_redis(listener, listener_ready_tx, responses).await;
280
0
    });
281
282
21
    listener_ready_rx
283
21
        .await
284
21
        .expect("Expected successful listener boot");
285
21
    info!(port, "Fake redis booted");
286
287
21
    port
288
21
}
289
290
21
pub async fn make_fake_redis_with_responses<B: BuildHasher + Clone + Send + 'static + Sync>(
291
21
    responses: HashMap<String, String, B>,
292
21
) -> u16 {
293
21
    make_fake_redis_with_multiple_responses(vec![responses]).await
294
21
}