Coverage Report

Created: 2026-06-04 10:48

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 tracing::{error, info, warn};
25
26
81
fn cmd_as_string(cmd: &redis::Cmd) -> String {
27
81
    let raw = cmd.get_packed_command();
28
81
    String::from_utf8(raw).unwrap()
29
81
}
30
31
890
pub(crate) fn arg_as_string(output: &mut String, arg: Value) {
32
890
    match arg {
33
98
        Value::SimpleString(s) => {
34
98
            write!(output, "+{s}\r\n").unwrap();
35
98
        }
36
114
        Value::Okay => {
37
114
            write!(output, "+OK\r\n").unwrap();
38
114
        }
39
200
        Value::BulkString(s) => {
40
200
            write!(
41
200
                output,
42
200
                "${}\r\n{}\r\n",
43
200
                s.len(),
44
200
                str::from_utf8(&s).unwrap()
45
200
            )
46
200
            .unwrap();
47
200
        }
48
190
        Value::Int(v) => {
49
190
            write!(output, ":{v}\r\n").unwrap();
50
190
        }
51
187
        Value::Array(values) => {
52
187
            write!(output, "*{}\r\n", values.len()).unwrap();
53
368
            for value in 
values187
{
54
368
                arg_as_string(output, value);
55
368
            }
56
        }
57
52
        Value::Map(values) => {
58
52
            write!(output, "%{}\r\n", values.len()).unwrap();
59
73
            for (key, value) in 
values52
{
60
73
                arg_as_string(output, key);
61
73
                arg_as_string(output, value);
62
73
            }
63
        }
64
48
        Value::Nil => {
65
48
            write!(output, "_\r\n").unwrap();
66
48
        }
67
1
        Value::Boolean(value) => {
68
1
            if value {
69
1
                write!(output, "#t\r\n")
70
            } else {
71
0
                write!(output, "#f\r\n")
72
            }
73
1
            .unwrap();
74
        }
75
        _ => {
76
0
            panic!("No support for {arg:?}")
77
        }
78
    }
79
890
}
80
81
65
fn args_as_string(args: Vec<Value>) -> String {
82
65
    let mut output = String::new();
83
163
    for arg in 
args65
{
84
163
        arg_as_string(&mut output, arg);
85
163
    }
86
65
    output
87
65
}
88
89
33
pub fn add_to_response<B: BuildHasher>(
90
33
    response: &mut HashMap<String, String, B>,
91
33
    cmd: &redis::Cmd,
92
33
    args: Vec<Value>,
93
33
) {
94
33
    add_to_response_raw(response, cmd, args_as_string(args));
95
33
}
96
97
33
pub fn add_to_response_raw<B: BuildHasher>(
98
33
    response: &mut HashMap<String, String, B>,
99
33
    cmd: &redis::Cmd,
100
33
    args: String,
101
33
) {
102
33
    response.insert(cmd_as_string(cmd), args);
103
33
}
104
105
16
fn setinfo(responses: &mut HashMap<String, String>) {
106
    // We do raw inserts of command here, because the library sends 3/4 commands in one go
107
    // They always start with HELLO, then optionally SELECT, so we use this to differentiate
108
16
    let hello = cmd_as_string(redis::cmd("HELLO").arg("3"));
109
16
    let setinfo = cmd_as_string(
110
16
        redis::cmd("CLIENT")
111
16
            .arg("SETINFO")
112
16
            .arg("LIB-NAME")
113
16
            .arg("redis-rs"),
114
    );
115
16
    responses.insert(
116
16
        [hello.clone(), setinfo.clone()].join(""),
117
16
        args_as_string(vec![
118
16
            Value::Map(vec![(
119
16
                Value::SimpleString("server".into()),
120
16
                Value::SimpleString("redis".into()),
121
16
            )]),
122
16
            Value::Okay,
123
16
            Value::Okay,
124
        ]),
125
    );
126
16
    responses.insert(
127
16
        [hello, cmd_as_string(redis::cmd("SELECT").arg(3)), setinfo].join(""),
128
16
        args_as_string(vec![
129
16
            Value::Map(vec![(
130
16
                Value::SimpleString("server".into()),
131
16
                Value::SimpleString("redis".into()),
132
16
            )]),
133
16
            Value::Okay,
134
16
            Value::Okay,
135
16
            Value::Okay,
136
        ]),
137
    );
138
16
}
139
140
9
pub fn add_lua_script<B: BuildHasher>(
141
9
    responses: &mut HashMap<String, String, B>,
142
9
    lua_script: &str,
143
9
    hash: &str,
144
9
) {
145
9
    add_to_response(
146
9
        responses,
147
9
        redis::cmd("SCRIPT").arg("LOAD").arg(lua_script),
148
9
        vec![hash.into_redis_value()],
149
    );
150
9
}
151
152
9
pub fn fake_redis_stream() -> HashMap<String, String> {
153
9
    let mut responses = HashMap::new();
154
9
    setinfo(&mut responses);
155
    // Does setinfo as well, so need to respond to all 3
156
9
    add_to_response(
157
9
        &mut responses,
158
9
        redis::cmd("SELECT").arg("3"),
159
9
        vec![Value::Okay, Value::Okay, Value::Okay],
160
    );
161
9
    responses
162
9
}
163
164
3
pub fn fake_redis_sentinel_master_stream() -> HashMap<String, String> {
165
3
    let mut response = fake_redis_stream();
166
3
    add_to_response(
167
3
        &mut response,
168
3
        &redis::cmd("ROLE"),
169
3
        vec![Value::Array(vec![
170
3
            "master".into_redis_value(),
171
3
            0.into_redis_value(),
172
3
            Value::Array(vec![]),
173
3
        ])],
174
    );
175
3
    response
176
3
}
177
178
7
pub fn fake_redis_sentinel_stream(master_name: &str, redis_port: u16) -> HashMap<String, String> {
179
7
    let mut response = HashMap::new();
180
7
    setinfo(&mut response);
181
182
    // Not a full "sentinel masters" response, but enough for redis-rs
183
7
    let resp: Vec<(Value, Value)> = vec![
184
7
        ("name".into_redis_value(), master_name.into_redis_value()),
185
7
        ("ip".into_redis_value(), "127.0.0.1".into_redis_value()),
186
7
        (
187
7
            "port".into_redis_value(),
188
7
            i64::from(redis_port).into_redis_value(),
189
7
        ),
190
7
        ("flags".into_redis_value(), "master".into_redis_value()),
191
    ];
192
193
7
    add_to_response(
194
7
        &mut response,
195
7
        redis::cmd("SENTINEL").arg("MASTERS"),
196
7
        vec![Value::Array(vec![Value::Map(resp)])],
197
    );
198
7
    response
199
7
}
200
201
23
pub(crate) async fn fake_redis_internal<H>(listener: TcpListener, handlers: Vec<H>)
202
23
where
203
23
    H: Fn(&[u8]) -> String + Send + Clone + 'static + Sync,
204
23
{
205
23
    let mut handler_iter = handlers.iter().cloned().cycle();
206
    loop {
207
59
        info!(
208
            "Waiting for connection on {}",
209
59
            listener.local_addr().unwrap()
210
        );
211
59
        let Ok((
mut stream36
, _)) = listener.accept().await else {
212
0
            error!("accept error");
213
0
            panic!("error");
214
        };
215
36
        info!("Accepted new connection");
216
36
        let local_handler = handler_iter.next().unwrap();
217
36
        background_spawn!("thread", async move {
218
            loop {
219
100k
                let mut buf = vec![0; 8192];
220
100k
                let 
res100k
= stream.read(&mut buf).await.
unwrap100k
();
221
100k
                if res != 0 {
222
363
                    let output = local_handler(&buf[..res]);
223
363
                    if !output.is_empty() {
224
230
                        stream.write_all(output.as_bytes()).await.unwrap();
225
133
                    }
226
100k
                }
227
            }
228
        });
229
    }
230
}
231
232
17
async fn fake_redis<B>(listener: TcpListener, all_responses: Vec<HashMap<String, String, B>>)
233
17
where
234
17
    B: BuildHasher + Clone + Send + 'static + Sync,
235
17
{
236
17
    let funcs = all_responses
237
17
        .iter()
238
17
        .map(|responses| {
239
17
            info!("Responses are: {:?}", responses);
240
17
            let values = responses.clone();
241
180
            move |buf: &[u8]| -> String {
242
180
                let str_buf = String::from_utf8_lossy(buf).into_owned();
243
635
                for (key, value) in 
&values180
{
244
635
                    if str_buf.starts_with(key) {
245
47
                        info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n"));
246
47
                        return value.clone();
247
588
                    }
248
                }
249
133
                warn!(
250
                    "Unknown command: {}",
251
133
                    str_buf.chars().take(1000).collect::<String>()
252
                );
253
133
                String::new()
254
180
            }
255
17
        })
256
17
        .collect();
257
17
    fake_redis_internal(listener, funcs).await;
258
0
}
259
260
17
async fn make_fake_redis_with_multiple_responses<B: BuildHasher + Clone + Send + 'static + Sync>(
261
17
    responses: Vec<HashMap<String, String, B>>,
262
17
) -> u16 {
263
17
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
264
17
    let port = listener.local_addr().unwrap().port();
265
17
    info!(port, "Fake redis booted");
266
267
17
    background_spawn!("listener", async move {
268
17
        fake_redis(listener, responses).await;
269
0
    });
270
271
17
    port
272
17
}
273
274
17
pub async fn make_fake_redis_with_responses<B: BuildHasher + Clone + Send + 'static + Sync>(
275
17
    responses: HashMap<String, String, B>,
276
17
) -> u16 {
277
17
    make_fake_redis_with_multiple_responses(vec![responses]).await
278
17
}