Skip to main content

tank_tests/
kv_storage.rs

1#![allow(unused_imports)]
2use std::collections::BTreeMap;
3use std::str::FromStr;
4use std::sync::{Arc, LazyLock};
5use tank::{Entity, Executor, expr};
6use tokio::sync::Mutex;
7use uuid::Uuid;
8
9use crate::silent_logs;
10
11static MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
12
13#[derive(Entity, PartialEq, Debug, Clone)]
14#[tank(schema = "testing", primary_key = (Self::key2, Self::key1))]
15pub struct KV {
16    pub key1: String,
17    pub key2: String,
18    pub string_val: String,
19    pub int_val: i64,
20    pub small_int: i16,
21    pub unsigned_val: u32,
22    pub float_val: f32,
23    pub double_val: f64,
24    pub boolean_val: bool,
25    pub opt_string: Option<String>,
26    pub opt_int: Option<i64>,
27    pub opt_float: Option<f64>,
28    pub opt_bool: Option<bool>,
29    pub uuid: Uuid,
30    pub opt_uuid: Option<Uuid>,
31    #[cfg(not(feature = "disable-arrays"))]
32    pub fixed_bytes: [u8; 16],
33    #[cfg(not(feature = "disable-arrays"))]
34    pub numbers: [i32; 4],
35    #[cfg(not(feature = "disable-lists"))]
36    pub tags: Vec<String>,
37    #[cfg(not(feature = "disable-lists"))]
38    pub scores: Vec<i32>,
39    #[cfg(not(feature = "disable-lists"))]
40    pub floats: Option<Vec<f64>>,
41    #[cfg(not(feature = "disable-lists"))]
42    pub shared_strings: Arc<Vec<String>>,
43    #[cfg(not(feature = "disable-maps"))]
44    pub metadata: BTreeMap<String, String>,
45    #[cfg(not(feature = "disable-maps"))]
46    pub counters: Option<BTreeMap<String, i64>>,
47}
48
49pub async fn kv_storage(executor: &mut impl Executor) {
50    let _lock = MUTEX.lock().await;
51
52    // Setup
53    silent_logs! {
54        // Silent logs for Valkey/Redis
55        KV::drop_table(executor, true, false)
56            .await
57            .expect("Failed to drop KV table");
58    }
59    KV::create_table(executor, false, true)
60        .await
61        .expect("Failed to create KV table");
62
63    // Insert values
64    let first = KV {
65        key1: "first".into(),
66        key2: "aa".into(),
67        string_val: "hello".into(),
68        int_val: 42,
69        small_int: -7,
70        unsigned_val: 99,
71        float_val: 3.14,
72        double_val: 9.87654321,
73        boolean_val: true,
74        opt_string: Some("optional".into()),
75        opt_int: None,
76        opt_float: Some(1.5),
77        opt_bool: None,
78        uuid: Uuid::from_str("c41af414-54cf-49cb-96c6-364e9c42f294").unwrap(),
79        opt_uuid: Some(Uuid::from_str("b1cacbcb-b3e0-4502-bbcd-d4ef014d189b").unwrap()),
80        #[cfg(not(feature = "disable-arrays"))]
81        fixed_bytes: *b"abcdefghijklmnop",
82        #[cfg(not(feature = "disable-arrays"))]
83        numbers: [1, 2, 3, 4],
84        #[cfg(not(feature = "disable-lists"))]
85        tags: vec!["kv".into(), "test".into(), "valkey".into()],
86        #[cfg(not(feature = "disable-lists"))]
87        scores: vec![10, 20, 30],
88        #[cfg(not(feature = "disable-lists"))]
89        floats: Some(vec![1.1, 2.2, 3.3]),
90        #[cfg(not(feature = "disable-lists"))]
91        shared_strings: Arc::new(vec!["a".into(), "b".into(), "c".into()]),
92        #[cfg(not(feature = "disable-maps"))]
93        metadata: BTreeMap::from_iter([
94            ("env".into(), "test".into()),
95            ("region".into(), "eu".into()),
96        ]),
97        #[cfg(not(feature = "disable-maps"))]
98        counters: Some(BTreeMap::from_iter([
99            ("views".into(), 100),
100            ("likes".into(), 5),
101        ])),
102    };
103    let second = KV {
104        key1: "second".into(),
105        key2: "bb".into(),
106        string_val: "world".into(),
107        int_val: 99,
108        small_int: 16,
109        unsigned_val: 0,
110        float_val: 0.005,
111        double_val: 0.00001,
112        boolean_val: false,
113        opt_string: None,
114        opt_int: 19.into(),
115        opt_float: None,
116        opt_bool: Some(true),
117        uuid: Uuid::from_str("4ba01a6a-e1b5-47c8-8d1b-e38b64a73989").unwrap(),
118        opt_uuid: None,
119        #[cfg(not(feature = "disable-arrays"))]
120        fixed_bytes: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
121        #[cfg(not(feature = "disable-arrays"))]
122        numbers: [-10, -20, -30, -40],
123        #[cfg(not(feature = "disable-lists"))]
124        tags: vec![],
125        #[cfg(not(feature = "disable-lists"))]
126        scores: vec![],
127        #[cfg(not(feature = "disable-lists"))]
128        floats: None,
129        #[cfg(not(feature = "disable-lists"))]
130        shared_strings: Arc::new(vec![]),
131        #[cfg(not(feature = "disable-maps"))]
132        metadata: BTreeMap::new(),
133        #[cfg(not(feature = "disable-maps"))]
134        counters: None,
135    };
136    let result = KV::insert_many(executor, [&first, &second])
137        .await
138        .expect("Failed to insert KV entity");
139
140    if let Some(rows) = result.rows_affected {
141        assert_eq!(rows, 2);
142    }
143
144    // Query (first, aa)
145    let value = KV::find_one(executor, expr!(KV::key1 == "first" && KV::key2 == "aa"))
146        .await
147        .expect("Failed to query KV");
148    assert_eq!(value, Some(first));
149
150    // Query (second, bb)
151    let value = KV::find_one(executor, expr!(KV::key1 == "second" && KV::key2 == "bb"))
152        .await
153        .expect("Failed to query KV");
154    assert_eq!(value, Some(second));
155}