mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
113 lines
6.3 KiB
Rust
113 lines
6.3 KiB
Rust
use litellm_cache::Error;
|
|
use litellm_cache_disk::{PythonDiskCacheAdapter, StoredValue, ValueAdapter};
|
|
use rstest::rstest;
|
|
|
|
enum ReadExpectation {
|
|
Bytes(&'static [u8]),
|
|
Miss,
|
|
Invalid,
|
|
}
|
|
|
|
#[rstest]
|
|
#[case::pickled_dictionary_with_string_keys(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e]),
|
|
ReadExpectation::Bytes(br#"{"a":1}"#)
|
|
)]
|
|
#[case::pickled_list_of_integers(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x65, 0x2e]),
|
|
ReadExpectation::Bytes(br#"[1,2]"#)
|
|
)]
|
|
#[case::pickled_tuple_of_integers(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x07, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4b, 0x01, 0x4b, 0x02, 0x86, 0x94, 0x2e]),
|
|
ReadExpectation::Bytes(br#"[1,2]"#)
|
|
)]
|
|
#[case::pickled_set_of_integers(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8f, 0x94, 0x28, 0x4b, 0x01, 0x4b, 0x02, 0x90, 0x2e]),
|
|
ReadExpectation::Bytes(br#"[1,2]"#)
|
|
)]
|
|
#[case::pickled_response_envelope(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x28, 0x8c, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x94, 0x47, 0x3f, 0xf8, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x72, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x94, 0x8c, 0x08, 0x7b, 0x22, 0x61, 0x22, 0x3a, 0x20, 0x31, 0x7d, 0x94, 0x75, 0x2e]),
|
|
ReadExpectation::Bytes(br#"{"response":"{\"a\": 1}","timestamp":1.5}"#)
|
|
)]
|
|
#[case::non_json_text(
|
|
StoredValue::Text("not json".into()),
|
|
ReadExpectation::Bytes(b"not json")
|
|
)]
|
|
#[case::json_text(
|
|
StoredValue::Text("{\"a\": 1}".into()),
|
|
ReadExpectation::Bytes(br#"{"a": 1}"#)
|
|
)]
|
|
#[case::non_utf8_bytes(
|
|
StoredValue::Bytes(vec![0xff, 0xfe]),
|
|
ReadExpectation::Bytes(&[0xff, 0xfe])
|
|
)]
|
|
#[case::integer_seven(StoredValue::Integer(7), ReadExpectation::Bytes(b"7"))]
|
|
#[case::float_one_point_five(StoredValue::Float(1.5), ReadExpectation::Bytes(b"1.5"))]
|
|
#[case::pickled_true(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e]),
|
|
ReadExpectation::Bytes(b"true")
|
|
)]
|
|
#[case::pickled_negative_integer(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x06, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4a, 0xfd, 0xff, 0xff, 0xff, 0x2e]),
|
|
ReadExpectation::Bytes(b"-3")
|
|
)]
|
|
#[case::pickled_bytes(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x09, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x43, 0x05, 0x62, 0x79, 0x74, 0x65, 0x73, 0x94, 0x2e]),
|
|
ReadExpectation::Invalid
|
|
)]
|
|
#[case::pickled_dictionary_with_integer_key(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x4b, 0x01, 0x8c, 0x01, 0x61, 0x94, 0x73, 0x2e]),
|
|
ReadExpectation::Invalid
|
|
)]
|
|
#[case::pickled_complex(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x2e, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x08, 0x62, 0x75, 0x69, 0x6c, 0x74, 0x69, 0x6e, 0x73, 0x94, 0x8c, 0x07, 0x63, 0x6f, 0x6d, 0x70, 0x6c, 0x65, 0x78, 0x94, 0x93, 0x94, 0x47, 0x3f, 0xf0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x86, 0x94, 0x52, 0x94, 0x2e]),
|
|
ReadExpectation::Invalid
|
|
)]
|
|
#[case::truncated_pickle(
|
|
StoredValue::Pickle(vec![0x80, 0x05, 0x2e]),
|
|
ReadExpectation::Invalid
|
|
)]
|
|
#[case::empty_bytes(StoredValue::Bytes(Vec::new()), ReadExpectation::Miss)]
|
|
#[case::empty_text(StoredValue::Text(String::new()), ReadExpectation::Miss)]
|
|
#[case::zero_integer(StoredValue::Integer(0), ReadExpectation::Miss)]
|
|
#[case::zero_float(StoredValue::Float(0.0), ReadExpectation::Miss)]
|
|
#[case::pickled_none(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_false(StoredValue::Pickle(vec![0x80, 0x05, 0x89, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_zero(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x00, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_zero_float(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x47, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_empty_string(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x8c, 0x00, 0x94, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_empty_list(StoredValue::Pickle(vec![0x80, 0x05, 0x5d, 0x94, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_empty_dictionary(StoredValue::Pickle(vec![0x80, 0x05, 0x7d, 0x94, 0x2e]), ReadExpectation::Miss)]
|
|
#[case::pickled_empty_tuple(StoredValue::Pickle(vec![0x80, 0x05, 0x29, 0x2e]), ReadExpectation::Miss)]
|
|
fn python_read_cases(#[case] row: StoredValue, #[case] expected: ReadExpectation) {
|
|
let result = PythonDiskCacheAdapter.read(row);
|
|
match expected {
|
|
ReadExpectation::Bytes(expected) => assert_eq!(result.unwrap().unwrap(), expected),
|
|
ReadExpectation::Miss => assert_eq!(result.unwrap(), None),
|
|
ReadExpectation::Invalid => assert!(matches!(result, Err(Error::InvalidEntry))),
|
|
}
|
|
}
|
|
|
|
#[rstest]
|
|
#[case::integer_two(Some(StoredValue::Integer(2)), 2.0)]
|
|
#[case::float_three_point_five(Some(StoredValue::Float(3.5)), 0.0)]
|
|
#[case::text_not_a_number(Some(StoredValue::Text("not a number".into())), 0.0)]
|
|
#[case::text_five(Some(StoredValue::Text("5".into())), 5.0)]
|
|
#[case::text_three_point_five(Some(StoredValue::Text("3.5".into())), 0.0)]
|
|
#[case::pickled_true(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x88, 0x2e])), 1.0)]
|
|
#[case::pickled_two(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4b, 0x02, 0x2e])), 2.0)]
|
|
#[case::pickled_dictionary(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x95, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x7d, 0x94, 0x8c, 0x01, 0x61, 0x94, 0x4b, 0x01, 0x73, 0x2e])), 0.0)]
|
|
#[case::missing(None, 0.0)]
|
|
#[case::pickled_none(Some(StoredValue::Pickle(vec![0x80, 0x05, 0x4e, 0x2e])), 0.0)]
|
|
fn python_counter_seed_cases(#[case] row: Option<StoredValue>, #[case] expected: f64) {
|
|
assert_eq!(PythonDiskCacheAdapter.counter_seed(row).unwrap(), expected);
|
|
}
|
|
|
|
#[rstest]
|
|
#[case::integer_three(3.0, StoredValue::Integer(3))]
|
|
#[case::fractional_three_point_five(3.5, StoredValue::Float(3.5))]
|
|
#[case::negative_zero(-0.0, StoredValue::Integer(0))]
|
|
#[case::large_float(1e300, StoredValue::Float(1e300))]
|
|
fn python_counter_value_cases(#[case] value: f64, #[case] expected: StoredValue) {
|
|
assert_eq!(PythonDiskCacheAdapter.counter_value(value), expected);
|
|
}
|