Skip to content

Commit c87ad91

Browse files
author
annie-mac
committed
add tests
1 parent 3f5dd25 commit c87ad91

File tree

3 files changed

+227
-0
lines changed

3 files changed

+227
-0
lines changed

‎sdk/cosmos/azure-cosmos-spark_3-5_2-12/src/test/scala/com/azure/cosmos/spark/ChangeFeedMetricsTest.scala

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
package com.azure.cosmos.spark
4+
5+
import com.azure.cosmos.changeFeedMetrics.ChangeFeedMetricsTracker
6+
7+
class ChangeFeedMetricsTest extends UnitSpec {
8+
9+
"ChangeFeedMetricsTracker" should "track weighted changes per lsn" in {
10+
val testRange = NormalizedRange("0", "FF")
11+
val metricsTracker = new ChangeFeedMetricsTracker(1, testRange)
12+
metricsTracker.track(10000, 0)
13+
metricsTracker.getWeightedAvgChangesPerLsn shouldBe 10000
14+
}
15+
16+
"ChangeFeedMetricsTracker" should "return none when no metrics tracked" in {
17+
val testRange = NormalizedRange("0", "FF")
18+
val metricsTracker = new ChangeFeedMetricsTracker(1, testRange)
19+
20+
metricsTracker.getWeightedAvgChangesPerLsn shouldBe None
21+
}
22+
23+
"ChangeFeedMetricsTracker" should "track limited metrics history" in {
24+
val testRange = NormalizedRange("0", "FF")
25+
val metricsTracker = new ChangeFeedMetricsTracker(1, testRange)
26+
27+
metricsTracker.track(10000, 1)
28+
for (i <- 1 to 5) {
29+
metricsTracker.track(1, 2000)
30+
}
31+
32+
metricsTracker.getWeightedAvgChangesPerLsn shouldBe 1.toDouble / 2000
33+
}
34+
}

‎sdk/cosmos/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/ChangeFeedMetricsTrackerSpec.scala

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
4+
package com.azure.cosmos.spark
5+
6+
import com.azure.cosmos.changeFeedMetrics.ChangeFeedMetricsTracker
7+
8+
class ChangeFeedMetricsTrackerSpec extends UnitSpec {
9+
10+
it should "track weighted changes per lsn" in {
11+
val testRange = NormalizedRange("0", "FF")
12+
val metricsTracker = new ChangeFeedMetricsTracker(1, testRange, 5, 0.5)
13+
// Single tracking
14+
metricsTracker.track(100, 200) // 2 changes per LSN
15+
metricsTracker.getWeightedAvgChangesPerLsn.get shouldBe 2.0
16+
17+
// Multiple tracking
18+
metricsTracker.track(50, 150) // 3 changes per LSN
19+
// With default decay factor 0.5, weighted avg should be (2.0 * 0.5 + 3.0 * 1.0)/(0.5 + 1.0) = 2.67
20+
metricsTracker.getWeightedAvgChangesPerLsn.get shouldBe 2.67 +- 0.01
21+
}
22+
23+
it should "return none when no metrics tracked" in {
24+
val testRange = NormalizedRange("0", "FF")
25+
val metricsTracker = new ChangeFeedMetricsTracker(1, testRange)
26+
27+
metricsTracker.getWeightedAvgChangesPerLsn shouldBe None
28+
}
29+
30+
it should "handle zero LSN gap correctly" in {
31+
val tracker = new ChangeFeedMetricsTracker(0L, NormalizedRange("", "FF"), 5, 0.5)
32+
33+
tracker.track(0, 100) // Should treat as 100 changes
34+
tracker.getWeightedAvgChangesPerLsn.get shouldBe 100.0
35+
36+
tracker.track(0, 0) // Should treat as 1 change
37+
tracker.getWeightedAvgChangesPerLsn.get shouldBe (100.0 * 0.5 + 1.0)/(0.5 + 1.0) +- 0.01
38+
}
39+
40+
41+
it should "respect maxHistory limit" in {
42+
val maxHistory = 2
43+
val tracker = new ChangeFeedMetricsTracker(0L, NormalizedRange("", "FF"), maxHistory, 0.5)
44+
45+
tracker.track(10, 10) // 1.0 changes per LSN
46+
tracker.track(10, 20) // 2.0 changes per LSN
47+
tracker.track(10, 30) // 3.0 changes per LSN - should evict first entry
48+
49+
// Should only consider last 2 entries: (2.0 * 0.5 + 3.0 * 1.0)/(0.5 + 1.0)
50+
tracker.getWeightedAvgChangesPerLsn.get shouldBe 2.67 +- 0.01
51+
}
52+
53+
it should "handle minimum change count of 1" in {
54+
val tracker = new ChangeFeedMetricsTracker(0L, NormalizedRange("", "FF"))
55+
56+
// Small values should be normalized to minimum of 1 change
57+
tracker.track(1000, 0)
58+
tracker.getWeightedAvgChangesPerLsn.get shouldBe 0.001 // 1/1000
59+
}
60+
}

‎sdk/cosmos/azure-cosmos-spark_3_2-12/src/test/scala/com/azure/cosmos/spark/CosmosPartitionPlannerSpec.scala

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33
package com.azure.cosmos.spark
44

55
import com.azure.core.management.AzureEnvironment
6+
import com.azure.cosmos.changeFeedMetrics.ChangeFeedMetricsTracker
67
import com.azure.cosmos.{ReadConsistencyStrategy, spark}
78
import org.apache.spark.sql.connector.read.streaming.ReadLimit
89

910
import java.time.Instant
1011
import java.util.UUID
12+
import java.util.concurrent.ConcurrentHashMap
1113
import java.util.concurrent.atomic.AtomicLong
1214

1315
class CosmosPartitionPlannerSpec extends UnitSpec {
@@ -484,4 +486,135 @@ class CosmosPartitionPlannerSpec extends UnitSpec {
484486
calculate(0).endLsn.get shouldEqual 2150
485487
calculate(1).endLsn.get shouldEqual 2150
486488
}
489+
490+
it should "calculateEndLsn should distribute rate based on metrics with readLimit and multiple partitions" in {
491+
val clientConfig = spark.CosmosClientConfiguration(
492+
UUID.randomUUID().toString,
493+
UUID.randomUUID().toString,
494+
CosmosMasterKeyAuthConfig(UUID.randomUUID().toString),
495+
None,
496+
UUID.randomUUID().toString,
497+
useGatewayMode = false,
498+
enforceNativeTransport = false,
499+
proactiveConnectionInitialization = None,
500+
proactiveConnectionInitializationDurationInSeconds = 120,
501+
httpConnectionPoolSize = 1000,
502+
readConsistencyStrategy = ReadConsistencyStrategy.EVENTUAL,
503+
disableTcpConnectionEndpointRediscovery = false,
504+
preferredRegionsList = Option.empty,
505+
subscriptionId = None,
506+
tenantId = None,
507+
resourceGroupName = None,
508+
azureEnvironmentEndpoints = AzureEnvironment.AZURE.getEndpoints,
509+
sparkEnvironmentInfo = "",
510+
clientBuilderInterceptors = None,
511+
clientInterceptors = None,
512+
sampledDiagnosticsLoggerConfig = None,
513+
azureMonitorConfig = None
514+
)
515+
516+
val containerConfig = CosmosContainerConfig(UUID.randomUUID().toString, UUID.randomUUID().toString)
517+
val normalizedRange = NormalizedRange(UUID.randomUUID().toString, UUID.randomUUID().toString)
518+
val docSizeInKB = rnd.nextInt()
519+
val maxRows = 10
520+
val nowEpochMs = Instant.now.toEpochMilli
521+
val createdAt = new AtomicLong(nowEpochMs)
522+
val lastRetrievedAt = new AtomicLong(nowEpochMs)
523+
524+
val metadata = PartitionMetadata(
525+
Map[String, String](),
526+
clientConfig,
527+
None,
528+
containerConfig,
529+
normalizedRange,
530+
documentCount = 2150,
531+
docSizeInKB,
532+
firstLsn = Some(0),
533+
latestLsn = 2150,
534+
startLsn = 2050,
535+
None,
536+
createdAt,
537+
lastRetrievedAt)
538+
539+
val metricsMap = new ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]()
540+
val metricsTracker = new ChangeFeedMetricsTracker(0L, normalizedRange)
541+
// Simulate metrics showing 2 changes per LSN on average
542+
metricsTracker.track(10, 20)
543+
metricsMap.put(normalizedRange, metricsTracker)
544+
545+
val calculate = CosmosPartitionPlanner.calculateEndLsn(
546+
Array[PartitionMetadata](metadata),
547+
ReadLimit.maxRows(maxRows),
548+
isChangeFeed = true,
549+
Some(metricsMap)
550+
)
551+
552+
// With 2 changes per LSN average from metrics and maxRows=10, should allow 5 LSN progress
553+
calculate(0).endLsn.get shouldEqual 2055
554+
}
555+
556+
it should "calculateEndLsn should handle when no progress is made even with metrics" in {
557+
val clientConfig = spark.CosmosClientConfiguration(
558+
UUID.randomUUID().toString,
559+
UUID.randomUUID().toString,
560+
CosmosMasterKeyAuthConfig(UUID.randomUUID().toString),
561+
None,
562+
UUID.randomUUID().toString,
563+
useGatewayMode = false,
564+
enforceNativeTransport = false,
565+
proactiveConnectionInitialization = None,
566+
proactiveConnectionInitializationDurationInSeconds = 120,
567+
httpConnectionPoolSize = 1000,
568+
readConsistencyStrategy = ReadConsistencyStrategy.EVENTUAL,
569+
disableTcpConnectionEndpointRediscovery = false,
570+
preferredRegionsList = Option.empty,
571+
subscriptionId = None,
572+
tenantId = None,
573+
resourceGroupName = None,
574+
azureEnvironmentEndpoints = AzureEnvironment.AZURE.getEndpoints,
575+
sparkEnvironmentInfo = "",
576+
clientBuilderInterceptors = None,
577+
clientInterceptors = None,
578+
sampledDiagnosticsLoggerConfig = None,
579+
azureMonitorConfig = None
580+
)
581+
582+
val containerConfig = CosmosContainerConfig(UUID.randomUUID().toString, UUID.randomUUID().toString)
583+
val normalizedRange = NormalizedRange(UUID.randomUUID().toString, UUID.randomUUID().toString)
584+
val docSizeInKB = rnd.nextInt()
585+
val maxRows = 10
586+
val nowEpochMs = Instant.now.toEpochMilli
587+
val createdAt = new AtomicLong(nowEpochMs)
588+
val lastRetrievedAt = new AtomicLong(nowEpochMs)
589+
590+
val metadata = PartitionMetadata(
591+
Map[String, String](),
592+
clientConfig,
593+
None,
594+
containerConfig,
595+
normalizedRange,
596+
documentCount = 2150,
597+
docSizeInKB,
598+
firstLsn = Some(0),
599+
latestLsn = 2050, // Latest LSN same as start LSN
600+
startLsn = 2050,
601+
None,
602+
createdAt,
603+
lastRetrievedAt)
604+
605+
val metricsMap = new ConcurrentHashMap[NormalizedRange, ChangeFeedMetricsTracker]()
606+
val metricsTracker = new ChangeFeedMetricsTracker(0L, normalizedRange)
607+
metricsTracker.track(2050, 100)
608+
metricsMap.put(normalizedRange, metricsTracker)
609+
610+
val calculate = CosmosPartitionPlanner.calculateEndLsn(
611+
Array[PartitionMetadata](metadata),
612+
ReadLimit.maxRows(maxRows),
613+
isChangeFeed = true,
614+
Some(metricsMap)
615+
)
616+
617+
// Should stay at start LSN since no progress can be made
618+
calculate(0).endLsn.get shouldEqual 2050
619+
}
487620
}

0 commit comments

Comments
 (0)