[Internal]Fixing test regressions introduced by recent PR adding new Otel attributes (Azure#45826)

* Fixing Otel attributes recently added

* Update CHANGELOG.md

* Update sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/TracerUnderTest.java

Co-authored-by: Copilot <[email protected]>

* Update sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Copilot <[email protected]>

Author: FabianMeiswinkel

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java

@@ -1481,7 +1481,26 @@ private void verifyOTelTracerAttributes(
 
             assertThat(attributes.get("azure.cosmosdb.operation.request_charge")).isNotNull();
             assertThat(attributes.get("azure.cosmosdb.operation.request_charge")).isInstanceOf(Double.class);
-            assertThat(attributes.get("azure.cosmosdb.operation.request_charge")).isEqualTo(Double.valueOf(ctx.getTotalRequestCharge()));
+
+            boolean isRequestChargeMismatch = 
+                (double)attributes.get("azure.cosmosdb.operation.request_charge") != 0d
+                || Double.valueOf(ctx.getTotalRequestCharge()) == 0d;
+
+            if (isRequestChargeMismatch) {
+                assertThat(attributes.get("azure.cosmosdb.operation.request_charge"))
+                    .isEqualTo(Double.valueOf(ctx.getTotalRequestCharge()));
+            } else {
+                // in the tests we usually use Flux.blockFirst() - which also cancels the Flux.
+                // Flux cancellation is also evident as an OTel event (which would have 0 request charge)
+                // to validate that previously the correct event - with right RU charge - had been traced we use the
+                // last (most recently) collected sibling Spans for the assertion.
+                assertThat(mockTracer.getAllCollectedSiblingSpans()).isNotNull();
+                TracerUnderTest.SpanRecord[] siblingSpans = mockTracer.getAllCollectedSiblingSpans();
+                assertThat(siblingSpans).hasAtLeastOneElementOfType(TracerUnderTest.SpanRecord.class);
+                TracerUnderTest.SpanRecord lastSiblingSpan = siblingSpans[siblingSpans.length - 1];
+                assertThat(lastSiblingSpan.getAttributes().get("azure.cosmosdb.operation.request_charge"))
+                    .isEqualTo(Double.valueOf(ctx.getTotalRequestCharge()));
+            }
 
             assertThat(attributes.get("azure.cosmosdb.response.sub_status_code")).isNotNull();
             assertThat(attributes.get("azure.cosmosdb.response.sub_status_code")).isInstanceOf(Integer.class);

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/TracerUnderTest.java

@@ -139,6 +139,18 @@ public synchronized void setAttribute(String key, String value, Context context)
         this.currentSpan.attributes.put(key, value);
     }
 
+    @Override
+    public synchronized void setAttribute(String key, long value, Context context) {
+        assertThat(this.currentSpan).isNotNull();
+        this.currentSpan.attributes.put(key, value);
+    }
+
+    @Override
+    public synchronized void setAttribute(String key, Object value, Context context) {
+        assertThat(this.currentSpan).isNotNull();
+        this.currentSpan.attributes.put(key, value);
+    }
+
     @Override
     public synchronized void addEvent(String name, Map<String, Object> attributes, OffsetDateTime timestamp, Context context) {
         assertThat(this.currentSpan).isNotNull();
@@ -156,6 +168,10 @@ public SpanRecord getCurrentSpan() {
         return this.currentSpan;
     }
 
+    public SpanRecord[] getAllCollectedSiblingSpans() {
+        return this.collectedSiblingSpans.toArray(new SpanRecord[0]);
+    }
+
     public Collection<EventRecord> getEventsOfAllCollectedSiblingSpans() {
         ArrayList<EventRecord> events = new ArrayList<>(this.currentSpan.events);
         LOGGER.info("getEventsOfAllCollectedSiblingSpans: {} {} - Siblings {}", this.currentSpan, this.currentSpan.events.size(), this.collectedSiblingSpans.size());

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -3,7 +3,7 @@
 ### 4.72.0-beta.1 (Unreleased)
 
 #### Features Added
-* Added `azure.cosmosdb.operation.request_charge` and `azure.cosmosdb.response.sub_status_code` trace attributes. - [PR 45753](https://github.com/Azure/azure-sdk-for-java/pull/45753)
+* Added `azure.cosmosdb.operation.request_charge` and `azure.cosmosdb.response.sub_status_code` trace attributes. - [PR 45753](https://github.com/Azure/azure-sdk-for-java/pull/45753) and [PR 45826](https://github.com/Azure/azure-sdk-for-java/pull/45826)
 
 #### Breaking Changes
 

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java

@@ -1540,7 +1540,7 @@ public void endSpan(CosmosDiagnosticsContext cosmosCtx, Context context, boolean
                 context);
             tracer.setAttribute(
                 "azure.cosmosdb.response.sub_status_code",
-                cosmosCtx.getSubStatusCode(),
+                (Integer)cosmosCtx.getSubStatusCode(),
                 context);
             tracer.setAttribute(
                 "db.cosmosdb.request_charge",