Skip to content

Commit 7e0e6ed

Browse files
ianbotsfianbotsf
authored
chore: Fix a flaky Kinesis test by always creating new resources (#1864)
* chore: Fix a flaky Kinesis test by always creating new resources * more careful resource handling
1 parent a782de3 commit 7e0e6ed

1 file changed

Lines changed: 122 additions & 133 deletions

File tree

‎services/kinesis/e2eTest/src/KinesisSubscribeToShardTest.kt‎

Lines changed: 122 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@ package aws.sdk.kotlin.services.kinesis
66

77
import aws.sdk.kotlin.services.kinesis.model.*
88
import aws.sdk.kotlin.services.kinesis.waiters.waitUntilStreamExists
9+
import aws.sdk.kotlin.services.kinesis.waiters.waitUntilStreamNotExists
910
import aws.sdk.kotlin.testing.withAllEngines
1011
import aws.smithy.kotlin.runtime.retries.getOrThrow
11-
import aws.smithy.kotlin.runtime.testing.AfterAll
12-
import aws.smithy.kotlin.runtime.testing.BeforeAll
1312
import aws.smithy.kotlin.runtime.testing.TestInstance
1413
import aws.smithy.kotlin.runtime.testing.TestLifecycle
14+
import aws.smithy.kotlin.runtime.time.Instant
15+
import aws.smithy.kotlin.runtime.time.TimestampFormat
1516
import kotlinx.coroutines.*
1617
import kotlinx.coroutines.flow.first
1718
import java.util.*
@@ -22,8 +23,8 @@ import kotlin.time.Duration.Companion.seconds
2223
private val WAIT_TIMEOUT = 30.seconds
2324
private val POLLING_RATE = 3.seconds
2425

25-
private val STREAM_NAME_PREFIX = "aws-sdk-kotlin-e2e-test-stream-"
26-
private val STREAM_CONSUMER_NAME_PREFIX = "aws-sdk-kotlin-e2e-test-"
26+
private val STREAM_NAME_PREFIX = "aws-sdk-kotlin-e2e-test-stream"
27+
private val STREAM_CONSUMER_NAME_PREFIX = "aws-sdk-kotlin-e2e-test"
2728

2829
private val TEST_DATA = "Bees, bees, bees, bees!"
2930

@@ -34,159 +35,147 @@ private val TEST_DATA = "Bees, bees, bees, bees!"
3435
class KinesisSubscribeToShardTest {
3536
private val client = KinesisClient { region = "us-east-1" }
3637

37-
private lateinit var dataStreamArn: String
38-
private lateinit var dataStreamConsumerArn: String
38+
private val testRunId = Instant
39+
.now()
40+
.format(TimestampFormat.ISO_8601_CONDENSED)
41+
.lowercase()
3942

40-
/**
41-
* Create infrastructure required for the test, if it doesn't exist already.
42-
*/
43-
@BeforeAll
44-
fun setup(): Unit = runBlocking {
45-
dataStreamArn = client.getOrCreateStream()
46-
dataStreamConsumerArn = client.getOrRegisterStreamConsumer()
47-
}
48-
49-
/**
50-
* Delete infrastructure used for the test.
51-
*/
52-
@AfterAll
53-
fun cleanUp(): Unit = runBlocking {
54-
client.deregisterStreamConsumer {
55-
streamArn = dataStreamArn
56-
consumerArn = dataStreamConsumerArn
57-
}
58-
59-
client.deleteStream {
60-
streamArn = dataStreamArn
61-
}
62-
}
43+
private val testStreamName = "$STREAM_NAME_PREFIX-$testRunId"
44+
private val testConsumerName = "$STREAM_CONSUMER_NAME_PREFIX-$testRunId"
6345

6446
/**
6547
* Select the single shard ID associated with the data stream, and subscribe to it.
6648
* Read one event and make sure the data matches what's expected.
6749
*/
6850
@Test
6951
fun testSubscribeToShard(): Unit = runBlocking {
70-
val dataStreamShardId = client.listShards {
71-
streamArn = dataStreamArn
72-
}.shards?.single()!!.shardId
73-
74-
withAllEngines { context ->
75-
client.withConfig {
76-
httpClient = context.engine
77-
}.use { clientWithTestEngine ->
78-
clientWithTestEngine.subscribeToShard(
79-
SubscribeToShardRequest {
80-
consumerArn = dataStreamConsumerArn
81-
shardId = dataStreamShardId
82-
startingPosition = StartingPosition {
83-
type = ShardIteratorType.TrimHorizon
52+
val dataStreamArn = client.createStream(testStreamName)
53+
54+
try {
55+
val dataStreamConsumerArn = client.registerStreamConsumerAndWait(testConsumerName, dataStreamArn)
56+
57+
try {
58+
val dataStreamShardId = client.listShards {
59+
streamArn = dataStreamArn
60+
}.shards?.single()!!.shardId
61+
62+
withAllEngines { context ->
63+
client.withConfig {
64+
httpClient = context.engine
65+
}.use { clientWithTestEngine ->
66+
clientWithTestEngine.subscribeToShard(
67+
SubscribeToShardRequest {
68+
consumerArn = dataStreamConsumerArn
69+
shardId = dataStreamShardId
70+
startingPosition = StartingPosition {
71+
type = ShardIteratorType.TrimHorizon
72+
}
73+
},
74+
) {
75+
val event = it.eventStream?.first()
76+
val record = event?.asSubscribeToShardEvent()?.records?.single()
77+
assertEquals(TEST_DATA, record?.data?.decodeToString())
8478
}
85-
},
86-
) {
87-
val event = it.eventStream?.first()
88-
val record = event?.asSubscribeToShardEvent()?.records?.single()
89-
assertEquals(TEST_DATA, record?.data?.decodeToString())
90-
}
9179

92-
// Wait 5 seconds, otherwise a ResourceInUseException gets thrown. Source:
93-
// https://docs.aws.amazon.com/kinesis/latest/APIReference/API_SubscribeToShard.html
94-
// > If you call SubscribeToShard 5 seconds or more after a successful call, the second call takes over the subscription
95-
delay(5.seconds)
80+
// Wait 5 seconds, otherwise a ResourceInUseException gets thrown. Source:
81+
// https://docs.aws.amazon.com/kinesis/latest/APIReference/API_SubscribeToShard.html
82+
// > If you call SubscribeToShard 5 seconds or more after a successful call, the second call takes over the subscription
83+
delay(5.seconds)
84+
}
85+
}
86+
} finally {
87+
client.deregisterStreamConsumer {
88+
streamArn = dataStreamArn
89+
consumerArn = dataStreamConsumerArn
90+
}
9691
}
92+
} finally {
93+
client.deleteStream { streamArn = dataStreamArn }
94+
client.waitUntilStreamNotExists { streamArn = dataStreamArn }
9795
}
9896
}
97+
}
9998

100-
/**
101-
* Get a Kinesis data stream with the [STREAM_NAME_PREFIX], or if one does not exist,
102-
* create one and populate it with one test record.
103-
* @return the ARN of the data stream
104-
*/
105-
private suspend fun KinesisClient.getOrCreateStream(): String = listStreams { }
106-
.streamSummaries
107-
?.find { it.streamName?.startsWith(STREAM_NAME_PREFIX) ?: false }
108-
?.streamArn ?: run {
109-
// Create a new data stream, then wait for it to be active
110-
val randomStreamName = STREAM_NAME_PREFIX + UUID.randomUUID()
111-
createStream {
112-
streamName = randomStreamName
113-
shardCount = 1
114-
}
99+
/**
100+
* Creates a Kinesis data stream with the prefix [STREAM_NAME_PREFIX] prefix and populate it with one test record.
101+
* @param name The name to use for the new stream
102+
* @return the ARN of the data stream
103+
*/
104+
private suspend fun KinesisClient.createStream(name: String): String {
105+
// Create a new data stream, then wait for it to be active
106+
createStream {
107+
streamName = name
108+
shardCount = 1
109+
}
115110

116-
val newStreamArn = waitUntilStreamExists({ streamName = randomStreamName })
117-
.getOrThrow()
118-
.streamDescription!!
119-
.streamArn!!
111+
val newStreamArn = waitUntilStreamExists { streamName = name }
112+
.getOrThrow()
113+
.streamDescription!!
114+
.streamArn
120115

121-
// Put a record, then wait for it to appear on the stream
122-
putRecord {
123-
data = TEST_DATA.encodeToByteArray()
124-
streamArn = newStreamArn
125-
partitionKey = "Goodbye"
126-
}
116+
// Put a record, then wait for it to appear on the stream
117+
putRecord {
118+
data = TEST_DATA.encodeToByteArray()
119+
streamArn = newStreamArn
120+
partitionKey = "Goodbye"
121+
}
127122

128-
val newStreamShardId = client.listShards {
129-
streamArn = newStreamArn
130-
}.shards?.single()!!.shardId
123+
val newStreamShardId = listShards {
124+
streamArn = newStreamArn
125+
}.shards?.single()!!.shardId
131126

132-
val currentShardIterator = getShardIterator {
133-
shardId = newStreamShardId
134-
shardIteratorType = ShardIteratorType.TrimHorizon
135-
streamArn = newStreamArn
136-
}.shardIterator!!
137-
138-
waitForResource {
139-
getRecords {
140-
shardIterator = currentShardIterator
141-
streamArn = newStreamArn
142-
}.records
143-
?.firstOrNull { it.data?.decodeToString() == TEST_DATA }
144-
}
127+
val currentShardIterator = getShardIterator {
128+
shardId = newStreamShardId
129+
shardIteratorType = ShardIteratorType.TrimHorizon
130+
streamArn = newStreamArn
131+
}.shardIterator!!
145132

146-
newStreamArn
133+
waitForResource {
134+
getRecords {
135+
shardIterator = currentShardIterator
136+
streamArn = newStreamArn
137+
}.records.firstOrNull { it.data.decodeToString() == TEST_DATA }
147138
}
148139

149-
/**
150-
* Get a Kinesis data stream consumer, or if it doesn't exist, register a new one.
151-
* @return the ARN of the stream consumer
152-
*/
153-
private suspend fun KinesisClient.getOrRegisterStreamConsumer(): String = listStreamConsumers { streamArn = dataStreamArn }
154-
.consumers
155-
?.firstOrNull { it.consumerName?.startsWith(STREAM_CONSUMER_NAME_PREFIX) ?: false }
156-
?.consumerArn ?: run {
157-
// Register a new consumer and wait for it to be active
158-
159-
val randomConsumerName = STREAM_CONSUMER_NAME_PREFIX + UUID.randomUUID()
160-
registerStreamConsumer {
161-
consumerName = randomConsumerName
162-
streamArn = dataStreamArn
163-
}
140+
return newStreamArn
141+
}
164142

165-
waitForResource {
166-
listStreamConsumers { streamArn = dataStreamArn }
167-
?.consumers
168-
?.firstOrNull { it.consumerName == randomConsumerName }
169-
?.takeIf { it.consumerStatus == ConsumerStatus.Active }
170-
?.consumerArn
171-
}
143+
/**
144+
* Register a new consumer and wait for it to be active
145+
* @param name The name to use for the new consumer
146+
* @param dataStreamArn The ARN of the data stream to consume
147+
* @return the ARN of the stream consumer
148+
*/
149+
private suspend fun KinesisClient.registerStreamConsumerAndWait(name: String, dataStreamArn: String): String {
150+
registerStreamConsumer {
151+
consumerName = name
152+
streamArn = dataStreamArn
172153
}
173154

174-
/**
175-
* Poll at a predefined [POLLING_RATE] for a resource to exist and return it.
176-
* Throws an exception if this takes longer than the [WAIT_TIMEOUT] duration.
177-
*
178-
* @param getResource a suspending function which returns the resource or null if it does not exist yet
179-
* @return the resource
180-
*/
181-
private suspend fun <T> KinesisClient.waitForResource(getResource: suspend () -> T?): T = withTimeout(WAIT_TIMEOUT) {
182-
var resource: T? = null
183-
while (resource == null) {
184-
resource = getResource()
185-
resource ?: run {
186-
delay(POLLING_RATE)
187-
yield()
188-
}
155+
return waitForResource {
156+
listStreamConsumers { streamArn = dataStreamArn }
157+
.consumers
158+
?.firstOrNull { it.consumerName == name }
159+
?.takeIf { it.consumerStatus == ConsumerStatus.Active }
160+
?.consumerArn
161+
}
162+
}
163+
164+
/**
165+
* Poll at a predefined [POLLING_RATE] for a resource to exist and return it.
166+
* Throws an exception if this takes longer than the [WAIT_TIMEOUT] duration.
167+
*
168+
* @param getResource a suspending function which returns the resource or null if it does not exist yet
169+
* @return the resource
170+
*/
171+
private suspend fun <T> waitForResource(getResource: suspend () -> T?): T = withTimeout(WAIT_TIMEOUT) {
172+
var resource: T? = null
173+
while (resource == null) {
174+
resource = getResource()
175+
resource ?: run {
176+
delay(POLLING_RATE)
177+
yield()
189178
}
190-
return@withTimeout resource
191179
}
180+
return@withTimeout resource
192181
}

0 commit comments

Comments
 (0)