From ab1f966a17939299324cbdb38178188f562c8880 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Thu, 17 Sep 2026 07:10:49 -0700 Subject: [PATCH] test coverage --- .../src/llms/cohere/ocr/transformation.rs | 91 ++++++++++----- .../src/llms/mistral/ocr/transformation.rs | 60 +++++----- .../crates/core/src/ocr/provider_config.rs | 7 ++ .../tests/azure_document_intelligence_ocr.rs | 105 +++++++++++++----- litellm-rust/crates/core/tests/reducto_ocr.rs | 72 ++++++++++-- 5 files changed, 247 insertions(+), 88 deletions(-) diff --git a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs index 09dd8d49757..996d9e462ab 100644 --- a/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/cohere/ocr/transformation.rs @@ -344,6 +344,7 @@ fn invalid_api_base() -> crate::ocr::Error { #[cfg(test)] mod tests { + use rstest::rstest; use serde_json::json; use super::*; @@ -471,10 +472,13 @@ mod tests { )); } - #[test] - fn provider_options_exclude_response_controls_and_extensions() { + #[rstest] + fn provider_options_exclude_response_controls_and_extensions( + #[values("markdown", "blocks")] output_format: &str, + #[values("https://example.com/a.png", "data:image/png;base64,YWJj")] source: &str, + ) { let arguments = serde_json::from_value( - json!({"output_format":"blocks","req_format":"native","unknown":true}), + json!({"output_format":output_format,"req_format":"native","unknown":true}), ) .unwrap(); let params = CohereParseConfig @@ -482,10 +486,10 @@ mod tests { .unwrap(); assert_eq!( serde_json::to_value(¶ms).unwrap(), - json!({"output_format":"blocks"}) + json!({"output_format":output_format}) ); let document = serde_json::from_value( - json!({"type":"image_url","image_url":"https://example.com/a.png","ignored":"field"}), + json!({"type":"image_url","image_url":source,"ignored":"field"}), ) .unwrap(); let body = CohereParseConfig @@ -494,7 +498,7 @@ mod tests { assert_eq!( serde_json::to_value(body).unwrap(), json!({ - "model":"parse", "document":{"type":"image_url","image_url":"https://example.com/a.png"}, "output_format":"blocks" + "model":"parse", "document":{"type":"image_url","image_url":source}, "output_format":output_format }) ); } @@ -526,9 +530,9 @@ mod tests { assert!(body.get("req_format").is_none()); } - #[test] + #[rstest] fn response_normalizes_markdown_images_blocks_and_billed_pages() { - let response = serde_json::from_value(json!({ + let payload = json!({ "pages": [ { "type":"markdown", @@ -558,17 +562,22 @@ mod tests { {"type":"blocks","blocks":[{"type":"text","text":{"content":"total"}}]} ], "meta":{"api_version":{"version":"2"},"billed_units":{"pages":3}} - })) - .unwrap(); + }); + let response = serde_json::from_value(payload.clone()).unwrap(); let normalized = normalize_response("parse-v5.0", response).unwrap(); assert_eq!(normalized.pages[0].index, 4); assert_eq!(normalized.pages[0].markdown, "receipt"); let image = &normalized.pages[0].images.as_ref().unwrap()[0]; - assert_eq!(image.bbox.as_ref().unwrap()["top_left_x"], 1); + let original_image = &payload["pages"][0]["markdown"]["images"][0]; assert_eq!( - image.extra_fields["bounding_box_normalized"]["bottom_right_x"], - 0.15 + serde_json::to_value(&image.bbox).unwrap(), + original_image["bounding_box"] ); + assert_eq!( + image.extra_fields["bounding_box_normalized"], + original_image["bounding_box_normalized"] + ); + assert_eq!(image.extra_fields["id"], original_image["id"]); assert_eq!(image.extra_fields["description"], "scan"); assert_eq!(image.extra_fields["category"], "logo"); assert_eq!(image.extra_fields["provider_extension"], "preserved"); @@ -609,9 +618,15 @@ mod tests { assert!(normalized.pages[0].images.is_none()); } - #[test] - fn response_types_documented_block_variants() { - let response = serde_json::from_value(json!({ + #[rstest] + fn response_types_documented_block_variants( + #[values( + crate::ocr::types::OcrResponseFormat::Litellm, + crate::ocr::types::OcrResponseFormat::Native + )] + response_format: crate::ocr::types::OcrResponseFormat, + ) { + let payload = json!({ "pages": [{ "type": "blocks", "index": 0, @@ -654,21 +669,45 @@ mod tests { "bottom_right_x": 0.7, "bottom_right_y": 0.8 }, - "title": "Totals" + "title": "Totals", + "description": "Invoice totals" } } ] }] - })) - .unwrap(); - let normalized = normalize_response("parse-v5.0", response).unwrap(); - let blocks = normalized.pages[0].extra_fields["blocks"] - .as_array() + }); + let normalized = CohereParseConfig + .transform_ocr_response( + "parse-v5.0", + &serde_json::to_vec(&payload).unwrap(), + response_format, + ) .unwrap(); - assert_eq!(blocks[0]["text"]["content"], "hello"); - assert_eq!(blocks[1]["image"]["category"], "logo"); - assert_eq!(blocks[2]["table"]["type"], "html"); - assert_eq!(blocks[2]["table"]["title"], "Totals"); + assert_eq!( + normalized.pages[0].extra_fields["blocks"], + payload["pages"][0]["blocks"] + ); + assert_eq!(normalized.pages[0].markdown, ""); + assert_eq!(normalized.pages[0].index, 0); + assert_eq!( + normalized.usage_info.as_ref().unwrap().pages_processed, + Some(1) + ); + match response_format { + crate::ocr::types::OcrResponseFormat::Litellm => { + assert!(normalized.provider_native_response.is_none()); + } + crate::ocr::types::OcrResponseFormat::Native => { + assert_eq!( + normalized.provider_native_response.as_ref(), + payload.as_object() + ); + } + } + assert_eq!( + normalized.into_json()["pages"][0]["blocks"], + payload["pages"][0]["blocks"] + ); } #[test] diff --git a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs index ffabce84d05..0f982e5e88a 100644 --- a/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/core/src/llms/mistral/ocr/transformation.rs @@ -425,7 +425,9 @@ mod tests { #[rstest] #[case("table_format", json!("html"))] + #[case("table_format", json!("markdown"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("page"))] #[case("document_annotation_prompt", json!("extract"))] #[case("include_blocks", json!(true))] #[case("id", json!("req-123"))] @@ -436,7 +438,9 @@ mod tests { #[rstest] #[case("pages", json!([0, 2]))] #[case("pages", json!("0,2-4"))] + #[case("pages", Value::Null)] #[case("include_image_base64", json!(true))] + #[case("include_image_base64", json!(false))] #[case("image_limit", json!(2))] #[case("image_min_size", json!(100))] #[case("bbox_annotation_format", json!({"type":"json_schema"}))] @@ -445,19 +449,28 @@ mod tests { #[case("extract_header", json!(true))] #[case("extract_footer", json!(false))] #[case("table_format", json!("html"))] + #[case("table_format", json!("markdown"))] #[case("confidence_scores_granularity", json!("word"))] + #[case("confidence_scores_granularity", json!("page"))] + #[case("confidence_scores_granularity", json!("block"))] #[case("include_blocks", json!(true))] + #[case("include_blocks", json!(false))] #[case("id", json!("req-123"))] - fn request_mapping_matches_python(#[case] name: &str, #[case] value: Value) { - let params: OpaqueParams = serde_json::from_value(json!({name: value.clone()})).unwrap(); + fn request_mapping_preserves_supplied_options(#[case] name: &str, #[case] value: Value) { + let arguments = serde_json::from_value(json!({name: value.clone()})).unwrap(); + let params = MistralOCRConfig + .map_ocr_params(&arguments, "model") + .unwrap(); let result = serde_json::to_value( MistralOCRConfig .transform_ocr_request("model", document(), ¶ms, &[]) .unwrap(), ) .unwrap(); - assert_eq!(result["model"], "model"); - assert_eq!(result[name], value); + assert_eq!( + result, + json!({"model":"model", "document":document(), name:value}) + ); } #[rstest] @@ -504,30 +517,25 @@ mod tests { #[rstest] fn transform_ocr_response_preserves_blocks_and_confidence_scores() { - let response: MistralOcrResponse = serde_json::from_value(json!({ - "pages":[{ - "index":0, - "markdown":"hello", - "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], - "dimensions":{"width":612,"height":792,"dpi":72}, - "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], - "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} - }], - "model":"returned-model", - "document_annotation":"{\"language\":\"en\"}", - "usage_info":{"pages_processed":1} - })) - .unwrap(); + let payload = json!({ + "pages":[{ + "index":0, + "markdown":"hello", + "images":[{"id":"img-0","image_base64":"data:image/png;base64,AA=="}], + "dimensions":{"width":612,"height":792,"dpi":72}, + "blocks":[{"type":"title","bbox":{"x":1},"confidence_scores":{"mean":0.98}}], + "confidence_scores":{"average_page_confidence_score":0.99,"minimum_page_confidence_score":0.97} + }], + "model":"returned-model", + "document_annotation":"{\"language\":\"en\"}", + "usage_info":{"pages_processed":1} + }); + let response: MistralOcrResponse = serde_json::from_value(payload.clone()).unwrap(); let result = normalize_response("model", response).unwrap().into_json(); - assert_eq!(result["pages"][0]["blocks"][0]["type"], "title"); - assert_eq!(result["pages"][0]["blocks"][0]["bbox"]["x"], 1); + assert_eq!(result["pages"][0]["blocks"], payload["pages"][0]["blocks"]); assert_eq!( - result["pages"][0]["blocks"][0]["confidence_scores"]["mean"], - 0.98 - ); - assert_eq!( - result["pages"][0]["confidence_scores"]["average_page_confidence_score"], - 0.99 + result["pages"][0]["confidence_scores"], + payload["pages"][0]["confidence_scores"] ); assert_eq!(result["pages"][0]["images"][0]["id"], "img-0"); assert_eq!(result["pages"][0]["dimensions"]["dpi"], 72); diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index ef9de23c913..37cc924fcc0 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -393,6 +393,13 @@ mod tests { #[rstest] #[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)] #[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)] + #[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/cohere-parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/cohere/parse-v5", OcrConfigKind::AzureCohere)] + #[case("azure_ai/invoice-parser", OcrConfigKind::AzureAi)] + #[case("azure_ai/parse-v5", OcrConfigKind::AzureAi)] + #[case("azure_ai/mistral-ocr-4-0", OcrConfigKind::AzureAi)] + #[case("azure_ai/mistral-document-ai-2512", OcrConfigKind::AzureAi)] #[case( "azure_ai/doc-intelligence/prebuilt-layout", OcrConfigKind::AzureDocumentIntelligence diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 1da340b57d4..41fe0c734cf 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,5 +1,6 @@ use std::sync::{Arc, Mutex}; +use rstest::rstest; use serde_json::{Value, json}; use super::test_support::{MockResponse, mock_server, perform_ocr, wire_request}; @@ -48,33 +49,87 @@ async fn facade_maps_pages_features_and_url_document() { ); } +#[rstest] +#[case(json!({"pages":[true]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages":[1,"2"]}), crate::ocr::Error::Pages("expected only integers or only strings".into()))] +#[case(json!({"pages":[-1]}), crate::ocr::Error::Pages("negative page index".into()))] +#[case(json!({"pages":"1&&features=bad"}), crate::ocr::Error::Pages("invalid native page range".into()))] +#[case(json!({"features":"languages&pages=1"}), crate::ocr::Error::Features)] +#[case(json!({"req_format":"azure"}), crate::ocr::Error::RequestFormat)] #[tokio::test] -async fn rejects_invalid_pages_features_and_format() { - for options in [ - json!({"pages":[true]}), - json!({"pages":[1,"2"]}), - json!({"pages":[-1]}), - json!({"pages":"1&&features=bad"}), - json!({"features":"languages&pages=1"}), - json!({"req_format":"azure"}), - ] { - let result = decode_request(OcrWireRequest { - model: "azure_ai/doc-intelligence/prebuilt-read".into(), - document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), - api_base: Some("http://127.0.0.1:1".into()), - custom_llm_provider: None, - extra_headers: None, - optional_params: options.as_object().unwrap().clone(), - input_sources: Default::default(), - timeout_seconds: None, - }); - let rejected = match result { - Ok(request) => perform_ocr(request).await.is_err(), - Err(_) => true, - }; - assert!(rejected, "accepted {options}"); +async fn rejects_invalid_pages_features_and_format( + #[case] options: Value, + #[case] expected: super::Error, +) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; + let result = decode_request(OcrWireRequest { + model: "azure_ai/doc-intelligence/prebuilt-read".into(), + document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), + api_key: Some("key".into()), + api_base: Some(base), + custom_llm_provider: None, + extra_headers: None, + optional_params: options.as_object().unwrap().clone(), + input_sources: Default::default(), + timeout_seconds: Some(2.0), + }); + let result = match result { + Ok(request) => perform_ocr(request).await, + Err(error) => Err(error), + }; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid options: {options}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); +} + +#[rstest] +#[case(json!({}))] +#[case(json!({"req_format":"litellm"}))] +#[tokio::test] +async fn missing_native_fields_keep_page_text_without_retaining_raw_response( + #[case] options: Value, +) { + let operation = json!({ + "status":"succeeded", + "analyzeResult":{"pages":[{"pageNumber":1,"lines":[{"content":"hello"}]}]} + }); + let (base, seen, server) = mock_server(vec![MockResponse::json(operation)]).await; + let response = perform_ocr(wire_request( + "azure_ai/doc-intelligence/prebuilt-read", + &base, + options, + )) + .await + .unwrap(); + server.await.unwrap(); + + assert_eq!(response.pages.len(), 1); + assert_eq!(response.pages[0].index, 0); + assert_eq!(response.pages[0].markdown, "hello"); + assert_eq!(response.provider_native_response, None); + let serialized = response.into_json(); + assert_eq!(serialized.get("content"), Some(&Value::Null)); + assert_eq!(serialized.get("tables"), Some(&Value::Null)); + assert_eq!(serialized.get("keyValuePairs"), Some(&Value::Null)); + let requests = seen.lock().unwrap(); + assert_eq!(requests.len(), 1); + let target = requests[0].split_whitespace().nth(1).unwrap(); + let url = format!("{base}{target}"); + for field in ["pages", "features", "req_format"] { + assert_eq!(query_value(&url, field), None); } + let body: Value = serde_json::from_str(requests[0].split_once("\r\n\r\n").unwrap().1).unwrap(); + assert_eq!(body, json!({"base64Source":"YWJj"})); } #[tokio::test] diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 0a7053b7429..0c25fd7a051 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -70,13 +70,26 @@ async fn request_mapping_matches_python( #[case("parse-v3")] #[case("parse-legacy")] #[tokio::test] -async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { +async fn data_uri_upload_preserves_multipart_headers( + #[case] model: &str, + #[values("application/pdf", "image/png")] mime_type: &str, +) { let (base, seen, server) = mock_server(vec![ MockResponse::json(json!({"file_id":"reducto://uploaded.pdf"})), MockResponse::json(json!({"result":{"chunks":[{"content":"hello"}]}})), ]) .await; - let mut request = wire_request(&format!("reducto/{model}"), &base, json!({})); + let document = if mime_type.starts_with("image/") { + json!({"type":"image_url","image_url":format!("data:{mime_type};base64,YWJj")}) + } else { + json!({"type":"document_url","document_url":format!("data:{mime_type};base64,YWJj")}) + }; + let mut request = super::LiteLLMOcrRequest { + document: serde_json::from_value::(document) + .unwrap() + .into(), + ..wire_request(&format!("reducto/{model}"), &base, json!({})) + }; request.transport.extra_headers = vec![ ("Content-Type".into(), "application/json".into()), ("X-Trace".into(), "upload-test".into()), @@ -94,9 +107,26 @@ async fn data_uri_upload_preserves_multipart_headers(#[case] model: &str) { .contains("content-type: multipart/form-data; boundary=") ); assert!(requests[0].contains("x-trace: upload-test")); - assert!(requests[0].contains("application/pdf")); - assert!(requests[0].contains("abc")); + let multipart = requests[0].split_once("\r\n\r\n").unwrap().1; + assert!(multipart.contains(&format!("Content-Type: {mime_type}\r\n"))); + assert!(multipart.contains("\r\n\r\nabc\r\n--")); assert!(requests[1].starts_with("POST /parse ")); + let source_field = if model == "parse-legacy" { + "document_url" + } else { + "input" + }; + assert_eq!( + request_body(&requests[1]), + json!({source_field:"reducto://uploaded.pdf"}) + ); + for request in requests.iter() { + assert!( + request + .to_ascii_lowercase() + .contains("authorization: bearer test-key\r\n") + ); + } } struct ParseBoundary { @@ -168,17 +198,37 @@ async fn upload_failure_stops_before_parse() { } #[rstest] -#[case("https://example.com/a.pdf")] -#[case("reducto://")] -#[case("data:application/pdf;base64")] -#[case("data:application/pdf;base64,INVALID!")] +#[case("https://example.com/a.pdf", crate::ocr::Error::ReductoSource)] +#[case("reducto://", crate::ocr::Error::RequestField { path: "document file id".into() })] +#[case("data:application/pdf;base64", crate::ocr::Error::InvalidDataUri)] +#[case( + "data:application/pdf;base64,INVALID!", + crate::ocr::Error::InvalidDataUri +)] #[tokio::test] -async fn rejects_invalid_document_sources_before_network(#[case] source: &str) { +async fn rejects_invalid_document_sources_before_network( + #[case] source: &str, + #[case] expected: super::Error, +) { + let (base, seen, server) = mock_server(vec![MockResponse::json(json!({}))]).await; let request = super::test_support::with_source( - wire_request("reducto/parse-v3", "http://127.0.0.1:1", json!({})), + wire_request("reducto/parse-v3", &base, json!({})), source, ); - assert!(perform_ocr(request).await.is_err()); + let result = perform_ocr(request).await; + server.abort(); + let _ = server.await; + assert!( + seen.lock().unwrap().is_empty(), + "sent invalid source: {source}" + ); + let error = result.unwrap_err(); + assert_eq!( + std::mem::discriminant(&error), + std::mem::discriminant(&expected) + ); + assert_eq!(error.http_status_code(), Some(400)); + assert_eq!(error.to_string(), expected.to_string()); } #[test]