From 6e16c6706e4ba54fb08250cb74e4ea2342d52cf7 Mon Sep 17 00:00:00 2001
From: Victoria Litvinova <73560279+vilit1@users.noreply.github.com>
Date: Fri, 6 Dec 2024 14:45:03 -0800
Subject: [PATCH] test: update the IoT Hub Cosmos Db custom endpoint tests to
 use the actual partition key (#726)

---
 azext_iot/tests/iothub/conftest.py                              | 1 +
 .../iothub/message_endpoint/test_iothub_message_endpoint_int.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/azext_iot/tests/iothub/conftest.py b/azext_iot/tests/iothub/conftest.py
index 81c719cbc..c5ffb0910 100644
--- a/azext_iot/tests/iothub/conftest.py
+++ b/azext_iot/tests/iothub/conftest.py
@@ -707,6 +707,7 @@ def _cosmos_db_provisioner():
         "cosmosdb": cosmos_obj,
         "database": database_obj,
         "container": container_obj,
+        "partitionKey": partition_key_path,
         "connectionString": _cosmos_db_get_cstring(account_name)
     }
 
diff --git a/azext_iot/tests/iothub/message_endpoint/test_iothub_message_endpoint_int.py b/azext_iot/tests/iothub/message_endpoint/test_iothub_message_endpoint_int.py
index c1e12b3e3..94de3486b 100644
--- a/azext_iot/tests/iothub/message_endpoint/test_iothub_message_endpoint_int.py
+++ b/azext_iot/tests/iothub/message_endpoint/test_iothub_message_endpoint_int.py
@@ -918,7 +918,7 @@ def test_iot_cosmos_endpoint_lifecycle(provisioned_cosmosdb_with_identity_module
     endpoint_names = generate_ep_names(3)
     partition_template = "{iothub}-{device_id}-{DD}-{MM}-{YYYY}"
     partition_template_default = "{deviceid}-{YYYY}-{MM}"
-    partition_path = "example"
+    partition_path = cosmosdb_obj["partitionKey"]
     # use connection string - no pkn or pkt
     cli.invoke(
         "iot hub message-endpoint create cosmosdb-container -n {} -g {} --en {} --erg {} -c {} --container {} "