From d8ce878a4142d2449d8658d4c7eb21ab9bea6357 Mon Sep 17 00:00:00 2001
From: Ryan Levick <me@ryanlevick.com>
Date: Fri, 30 Jul 2021 17:42:16 +0200
Subject: [PATCH] Fix issues discovered in e2e tests

---
 sdk/cosmos/src/clients/collection_client.rs  | 23 ++++++++++++++------
 sdk/cosmos/src/operations/create_document.rs |  4 ++--
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/sdk/cosmos/src/clients/collection_client.rs b/sdk/cosmos/src/clients/collection_client.rs
index af8c40d5d6..2c74ae6d21 100644
--- a/sdk/cosmos/src/clients/collection_client.rs
+++ b/sdk/cosmos/src/clients/collection_client.rs
@@ -4,9 +4,10 @@ use crate::clients::*;
 use crate::operations::*;
 use crate::requests;
 use crate::resources::ResourceType;
+use crate::CosmosEntity;
 use crate::ReadonlyString;
 use azure_core::PipelineContext;
-use azure_core::{pipeline::Pipeline, Context, HttpClient};
+use azure_core::{pipeline::Pipeline, Context, HttpClient, Request};
 use serde::Serialize;
 
 /// A client for Cosmos collection resources.
@@ -63,16 +64,14 @@ impl CollectionClient {
     }
 
     /// create a document in a collection
-    pub async fn create_document<D: Serialize>(
+    pub async fn create_document<'a, D: Serialize + CosmosEntity<'a>>(
         &self,
         ctx: Context,
-        document: &D,
+        document: &'a D,
         options: CreateDocumentOptions<'_>,
     ) -> Result<CreateDocumentResponse, crate::Error> {
-        let mut request = self
-            .cosmos_client()
-            .prepare_request_pipeline("collection", http::Method::POST);
-        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Collections.into());
+        let mut request = self.prepare_request_pipeline(http::Method::POST);
+        let mut pipeline_context = PipelineContext::new(ctx, ResourceType::Documents.into());
 
         options.decorate_request(&mut request, document)?;
         let response = self
@@ -162,4 +161,14 @@ impl CollectionClient {
     pub(crate) fn pipeline(&self) -> &Pipeline<CosmosContext> {
         self.cosmos_client().pipeline()
     }
+
+    fn prepare_request_pipeline(&self, http_method: http::Method) -> Request {
+        let path = &format!(
+            "dbs/{}/colls/{}",
+            self.database_client().database_name(),
+            self.collection_name()
+        );
+        self.cosmos_client()
+            .prepare_request_pipeline(&path, http_method)
+    }
 }
diff --git a/sdk/cosmos/src/operations/create_document.rs b/sdk/cosmos/src/operations/create_document.rs
index 4e3298312c..00ea4c3d7a 100644
--- a/sdk/cosmos/src/operations/create_document.rs
+++ b/sdk/cosmos/src/operations/create_document.rs
@@ -59,13 +59,13 @@ impl<'a> CreateDocumentOptions<'a> {
         document: &'b DOC,
     ) -> Result<(), crate::Error>
     where
-        DOC: Serialize,
+        DOC: Serialize + CosmosEntity<'b>,
     {
         let serialized = serde_json::to_string(document)?;
         let partition_key = self
             .partition_key
             .clone()
-            .unwrap_or_else(|| serialize_partition_key(document).unwrap());
+            .unwrap_or_else(|| serialize_partition_key(&document.partition_key()).unwrap());
 
         add_as_partition_key_header_serialized2(&partition_key, req);
         azure_core::headers::add_optional_header2(&self.if_match_condition, req)?;