[Cosmos] Integration tests for items and query

sdk/core/azure_core_test/src/recorded.rs

@@ -18,7 +18,14 @@ static TEST_PROXY: OnceCell<Result<Arc<Proxy>>> = OnceCell::const_new();
 ///
 /// The [Test Proxy](https://github.com/Azure/azure-sdk-tools/blob/main/tools/test-proxy/Azure.Sdk.Tools.TestProxy/README.md) service will be started as needed.
 /// Every `#[recorded::test]` will call this automatically, but it can also be called manually by any other test e.g., those attributed with `#[tokio::test]`.
-pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<Session> {
+///
+/// This function will return `Ok(None)` if the test is running in live mode and should not use the test proxy at all.
+pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<Option<Session>> {
+    // Live tests don't use test-proxy.
+    if ctx.test_mode() == azure_core::test::TestMode::Live {
+        return Ok(None);
+    }
+
     let proxy = TEST_PROXY
         .get_or_init(|| async move {
             proxy::start(ctx.test_data_dir(), options)
@@ -30,8 +37,8 @@ pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<S
         .map_err(|err| azure_core::Error::new(err.kind().clone(), err))?;
 
     let span = debug_span!(target: crate::SPAN_TARGET, "session", mode = ?ctx.test_mode(), test = ?ctx.test_name());
-    Ok(Session {
+    Ok(Some(Session {
         proxy: proxy.clone(),
         span,
-    })
+    }))
 }

sdk/cosmos/azure_data_cosmos/examples/cosmos/create.rs

@@ -2,7 +2,7 @@ use std::error::Error;
 
 use azure_data_cosmos::{
     models::{ContainerProperties, PartitionKeyDefinition, ThroughputProperties},
-    CosmosClient, CreateContainerOptions, CreateDatabaseOptions, PartitionKey,
+    CosmosClient, CreateContainerOptions, CreateDatabaseOptions, ItemOptions, PartitionKey,
 };
 use clap::{Args, Subcommand};
 
@@ -32,6 +32,10 @@ pub enum Subcommands {
         /// The JSON of the new item.
         #[clap(long, short)]
         json: String,
+
+        /// If set, the updated item will be included in the response.
+        #[clap(long)]
+        show_updated: bool,
     },
 
     /// Create a database (does not support Entra ID).
@@ -73,21 +77,30 @@ impl CreateCommand {
                 container,
                 partition_key,
                 json,
+                show_updated,
             } => {
                 let db_client = client.database_client(&database);
                 let container_client = db_client.container_client(&container);
 
                 let pk = PartitionKey::from(&partition_key);
                 let item: serde_json::Value = serde_json::from_str(&json)?;
 
-                let created = container_client
-                    .create_item(pk, item, None)
-                    .await?
-                    .into_body()
-                    .await?
-                    .unwrap();
-                println!("Created item:");
-                println!("{:#?}", created);
+                let options = ItemOptions {
+                    enable_content_response_on_write: show_updated,
+                    ..Default::default()
+                };
+
+                let response = container_client
+                    .create_item(pk, item, Some(options))
+                    .await?;
+
+                println!("Created item successfully");
+
+                if show_updated {
+                    let created = response.into_json_body::<serde_json::Value>().await?;
+                    println!("Newly created item:");
+                    println!("{:#?}", created);
+                }
                 Ok(())
             }
 
@@ -106,8 +119,7 @@ impl CreateCommand {
                     .create_database(&id, options)
                     .await?
                     .into_body()
-                    .await?
-                    .unwrap();
+                    .await?;
                 println!("Created database:");
                 println!("{:#?}", db);
                 Ok(())
@@ -152,8 +164,7 @@ impl CreateCommand {
                     .create_container(properties, options)
                     .await?
                     .into_body()
-                    .await?
-                    .unwrap();
+                    .await?;
                 println!("Created container:");
                 println!("{:#?}", container);
                 Ok(())

sdk/cosmos/azure_data_cosmos/examples/cosmos/read.rs

@@ -61,7 +61,7 @@ impl ReadCommand {
                         println!("Item not found!")
                     }
                     Ok(r) => {
-                        let item: serde_json::Value = r.into_body().await?.unwrap();
+                        let item: serde_json::Value = r.into_json_body().await?;
                         println!("Found item:");
                         println!("{:#?}", item);
                     }

sdk/cosmos/azure_data_cosmos/examples/cosmos/replace.rs

@@ -1,7 +1,7 @@
 use std::error::Error;
 
 use azure_core::StatusCode;
-use azure_data_cosmos::{CosmosClient, PartitionKey};
+use azure_data_cosmos::{CosmosClient, ItemOptions, PartitionKey};
 use clap::{Args, Subcommand};
 
 use crate::utils::ThroughputOptions;
@@ -33,6 +33,10 @@ pub enum Subcommands {
         /// The JSON of the new item.
         #[clap(long, short)]
         json: String,
+
+        /// If set, the updated item will be included in the response.
+        #[clap(long)]
+        show_updated: bool,
     },
     DatabaseThroughput {
         /// The database to update throughput for.
@@ -62,24 +66,34 @@ impl ReplaceCommand {
                 item_id,
                 partition_key,
                 json,
+                show_updated,
             } => {
                 let db_client = client.database_client(&database);
                 let container_client = db_client.container_client(&container);
 
                 let pk = PartitionKey::from(&partition_key);
                 let item: serde_json::Value = serde_json::from_str(&json)?;
 
+                let options = ItemOptions {
+                    enable_content_response_on_write: show_updated,
+                    ..Default::default()
+                };
+
                 let response = container_client
-                    .replace_item(pk, &item_id, item, None)
+                    .replace_item(pk, &item_id, item, Some(options))
                     .await;
                 match response {
                     Err(e) if e.http_status() == Some(StatusCode::NotFound) => {
                         println!("Item not found!")
                     }
                     Ok(r) => {
-                        let item: serde_json::Value = r.into_body().await?.unwrap();
-                        println!("Replaced item:");
-                        println!("{:#?}", item);
+                        println!("Replaced item successfully");
+
+                        if show_updated {
+                            let created: serde_json::Value = r.into_json_body().await?;
+                            println!("Newly replaced item:");
+                            println!("{:#?}", created);
+                        }
                     }
                     Err(e) => return Err(e.into()),
                 };

sdk/cosmos/azure_data_cosmos/examples/cosmos/upsert.rs

@@ -1,6 +1,6 @@
 use std::error::Error;
 
-use azure_data_cosmos::{CosmosClient, PartitionKey};
+use azure_data_cosmos::{CosmosClient, ItemOptions, PartitionKey};
 use clap::Args;
 
 /// Creates a new item or replaces an existing item, if a matching item already exists.
@@ -19,6 +19,10 @@ pub struct UpsertCommand {
     /// The JSON of the new item.
     #[clap(long, short)]
     json: String,
+
+    /// If set, the updated item will be included in the response.
+    #[clap(long)]
+    show_updated: bool,
 }
 
 impl UpsertCommand {
@@ -29,14 +33,21 @@ impl UpsertCommand {
         let pk = PartitionKey::from(&self.partition_key);
         let item: serde_json::Value = serde_json::from_str(&self.json)?;
 
-        let created = container_client
-            .upsert_item(pk, item, None)
-            .await?
-            .into_body()
-            .await?
-            .unwrap();
-        println!("Created item:");
-        println!("{:#?}", created);
+        let options = ItemOptions {
+            enable_content_response_on_write: self.show_updated,
+            ..Default::default()
+        };
+
+        let response = container_client
+            .upsert_item(pk, item, Some(options))
+            .await?;
+        println!("Item updated successfully");
+
+        if self.show_updated {
+            let created: serde_json::Value = response.into_json_body().await?;
+            println!("Updated item:");
+            println!("{:#?}", created);
+        }
         Ok(())
     }
 }