rabbitmq_stream_client/
byte_capacity.rs

1pub const KILOBYTE: u64 = 1000;
2pub const MEGABYTE: u64 = 1000 * KILOBYTE;
3pub const GIGABYTE: u64 = 1000 * MEGABYTE;
4pub const TERABYTE: u64 = 1000 * GIGABYTE;
5
6pub enum ByteCapacity {
7    B(u64),
8    KB(u64),
9    MB(u64),
10    GB(u64),
11    TB(u64),
12}
13
14impl ByteCapacity {
15    pub fn bytes(&self) -> u64 {
16        match self {
17            ByteCapacity::B(x) => *x,
18            ByteCapacity::KB(x) => x * KILOBYTE,
19            ByteCapacity::MB(x) => x * MEGABYTE,
20            ByteCapacity::GB(x) => x * GIGABYTE,
21            ByteCapacity::TB(x) => x * TERABYTE,
22        }
23    }
24}
25
26#[cfg(feature = "serde")]
27impl<'de> serde::Deserialize<'de> for ByteCapacity {
28    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
29    where
30        D: serde::Deserializer<'de>,
31    {
32        use serde::Deserialize;
33
34        #[derive(Deserialize)]
35        #[serde(untagged)]
36        enum StringOrNumber {
37            String(String),
38            Number(u64),
39        }
40
41        macro_rules! match_suffix {
42            ($str: ident, $suf: expr, $variant: expr) => {
43                if $str.ends_with($suf) {
44                    let num = $str
45                        .trim_end_matches($suf)
46                        .parse()
47                        .map_err(serde::de::Error::custom)?;
48                    return Ok(($variant)(num));
49                }
50            };
51        }
52
53        let s_or_n = StringOrNumber::deserialize(deserializer)?;
54
55        match s_or_n {
56            StringOrNumber::String(str) => {
57                match_suffix!(str, "TB", ByteCapacity::TB);
58                match_suffix!(str, "GB", ByteCapacity::GB);
59                match_suffix!(str, "MB", ByteCapacity::MB);
60                match_suffix!(str, "KB", ByteCapacity::KB);
61                match_suffix!(str, "B", ByteCapacity::B);
62
63                let num = str.parse().map_err(|_| {
64                    serde::de::Error::custom(
65                        "Expect a number or a string with a TB|GB|MB|KB|B suffix",
66                    )
67                })?;
68                Ok(ByteCapacity::B(num))
69            }
70            StringOrNumber::Number(num) => Ok(ByteCapacity::B(num)),
71        }
72    }
73}
74
75#[cfg(feature = "serde")]
76mod tests {
77    #[test]
78    fn test_deserilize_byte_capacity() {
79        use crate::types::ByteCapacity;
80
81        assert!(matches!(
82            serde_json::from_str::<ByteCapacity>("\"5GB\""),
83            Ok(ByteCapacity::GB(5))
84        ));
85        assert!(matches!(
86            serde_json::from_str::<ByteCapacity>("\"5TB\""),
87            Ok(ByteCapacity::TB(5))
88        ));
89        assert!(matches!(
90            serde_json::from_str::<ByteCapacity>("\"5MB\""),
91            Ok(ByteCapacity::MB(5))
92        ));
93        assert!(matches!(
94            serde_json::from_str::<ByteCapacity>("\"5KB\""),
95            Ok(ByteCapacity::KB(5))
96        ));
97        assert!(matches!(
98            serde_json::from_str::<ByteCapacity>("\"5B\""),
99            Ok(ByteCapacity::B(5))
100        ));
101        assert!(matches!(
102            serde_json::from_str::<ByteCapacity>("\"5\""),
103            Ok(ByteCapacity::B(5))
104        ));
105        assert!(matches!(
106            serde_json::from_str::<ByteCapacity>("5"),
107            Ok(ByteCapacity::B(5))
108        ));
109        let err = serde_json::from_str::<ByteCapacity>("\"Wrong string format\"")
110            .err()
111            .expect("Expect an error");
112        assert_eq!(
113            err.to_string(),
114            "Expect a number or a string with a TB|GB|MB|KB|B suffix"
115        );
116    }
117}
rabbitmq_stream_client/byte_capacity.rs

rabbitmq_stream_client/
byte_capacity.rs
