Merge pull request #698 from ButterflyNetwork/kinesis-error

[Kinesis] Ensure submitAllRecords returns an error on failure to submit

Author: minbi

AWSKinesis/AWSAbstractKinesisRecorder.m

@@ -279,12 +279,19 @@ - (AWSTask *)submitAllRecords {
 
                     NSString *streamName = temporaryRecords[0][@"stream_name"];
 
-                    [[self.recorderHelper submitRecordsForStream:streamName
-                                                         records:temporaryRecords
-                                                   partitionKeys:partitionKeys
-                                                putPartitionKeys:putPartitionKeys
-                                              retryPartitionKeys:retryPartitionKeys
-                                                            stop:&stop] waitUntilFinished];
+                    AWSTask *submitTask = \
+                        [self.recorderHelper submitRecordsForStream:streamName
+                                                            records:temporaryRecords
+                                                      partitionKeys:partitionKeys
+                                                   putPartitionKeys:putPartitionKeys
+                                                 retryPartitionKeys:retryPartitionKeys
+                                                               stop:&stop];
+
+                    [submitTask waitUntilFinished];
+
+                    if (submitTask.error) {
+                        error = submitTask.error;
+                    }
 
                     for (NSString *partitionKey in putPartitionKeys) {
                         BOOL result = [db executeUpdate:@"DELETE FROM record WHERE partition_key = :partition_key"

AWSKinesis/AWSFirehoseRecorder.m

@@ -177,8 +177,11 @@ - (AWSTask *)submitRecordsForStream:(NSString *)streamName
     return [[self.firehose putRecordBatch:putRecordBatchInput] continueWithBlock:^id(AWSTask *task) {
         if (task.error) {
             AWSDDLogError(@"Error: [%@]", task.error);
-            if ([task.error.domain isEqualToString:NSURLErrorDomain]) {
+            const NSArray *stopErrorDomains = @[NSURLErrorDomain, AWSCognitoIdentityErrorDomain];
+
+            if (task.error && [stopErrorDomains containsObject:task.error.domain]) {
                 *stop = YES;
+                return [AWSTask taskWithError:task.error];
             }
         }
         if (task.result) {

AWSKinesis/AWSKinesisRecorder.m

@@ -181,6 +181,8 @@ - (AWSTask *)submitRecordsForStream:(NSString *)streamName
             if ([task.error.domain isEqualToString:NSURLErrorDomain]) {
                 *stop = YES;
             }
+
+            return [AWSTask taskWithError:task.error];
         }
         if (task.result) {
             AWSKinesisPutRecordsOutput *putRecordsOutput = task.result;

AWSKinesisTests/AWSFirehoseRecorderTests.m

@@ -191,23 +191,35 @@ - (void)testDiskAgeLimit {
     FirehoseRecorder.diskAgeLimit = 0.0;
 }
 
-- (void)testAll {
-    AWSFirehoseRecorder *firehoseRecorder = [AWSFirehoseRecorder defaultFirehoseRecorder];
-    
-    NSMutableArray *tasks = [NSMutableArray new];
-    for (int32_t i = 0; i < 1234; i++) {
-        [tasks addObject:[firehoseRecorder saveRecord:[[NSString stringWithFormat:@"TestString-%02d\n", i] dataUsingEncoding:NSUTF8StringEncoding]
-                                           streamName:AWSFirehoseRecorderTestStream]];
-    }
-    
-    [[[[AWSTask taskForCompletionOfAllTasks:tasks] continueWithSuccessBlock:^id(AWSTask *task) {
-        sleep(10);
-        return [firehoseRecorder submitAllRecords];
-    }] continueWithBlock:^id(AWSTask *task) {
-        XCTAssertNil(task.error);
-        
+- (void)testSubmitAllRecordsReturnsErrorOnInvalidPoolId {
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Test finished running."];
+
+    NSString *poolId = @"invalidPoolId";
+    AWSCognitoCredentialsProvider *invalidCreds = \
+        [[AWSCognitoCredentialsProvider alloc] initWithRegionType:AWSRegionUSEast1
+                                                   identityPoolId:poolId];
+
+    AWSServiceConfiguration *configuration = \
+        [[AWSServiceConfiguration alloc] initWithRegion:AWSRegionUSEast1
+                                    credentialsProvider:invalidCreds];
+
+    [AWSFirehoseRecorder registerFirehoseRecorderWithConfiguration:configuration
+                                                            forKey:poolId];
+    AWSFirehoseRecorder *firehoseRecorder = [AWSFirehoseRecorder FirehoseRecorderForKey:poolId];
+    [firehoseRecorder saveRecord:[@"testString" dataUsingEncoding:NSUTF8StringEncoding]
+                      streamName:@"streamName"];
+
+    AWSTask *submitTask = firehoseRecorder.submitAllRecords;
+
+    [submitTask continueWithBlock:^id(AWSTask *task) {
+        XCTAssertNotNil(task.error, @"Task should have an error due to invalid pool id.");
+        [expectation fulfill];
         return nil;
-    }] waitUntilFinished];
+    }];
+
+    [self waitForExpectationsWithTimeout:5 handler:^(NSError * _Nullable error) {
+        XCTAssertNil(error);
+    }];
 }
 
 @end

AWSKinesisTests/AWSKinesisRecorderTests.m

@@ -397,6 +397,38 @@ - (void)testAllWithPartitionKey {
     }];
 }
 
+- (void)testSubmitAllRecordsReturnsError {
+    XCTestExpectation *expectation = [self expectationWithDescription:@"Test finished running."];
+
+    NSString *poolId = @"invalidPoolId";
+    AWSCognitoCredentialsProvider *invalidCreds = \
+        [[AWSCognitoCredentialsProvider alloc] initWithRegionType:AWSRegionUSEast1
+                                                   identityPoolId:poolId];
+
+    AWSServiceConfiguration *configuration = \
+        [[AWSServiceConfiguration alloc] initWithRegion:AWSRegionUSEast1
+                                    credentialsProvider:invalidCreds];
+
+    [AWSKinesisRecorder registerKinesisRecorderWithConfiguration:configuration
+                                                          forKey:poolId];
+
+    AWSKinesisRecorder *kinesisRecorder = [AWSKinesisRecorder KinesisRecorderForKey:poolId];
+    [kinesisRecorder saveRecord:[@"testString" dataUsingEncoding:NSUTF8StringEncoding]
+                     streamName:testStreamName];
+
+    AWSTask *submitTask = kinesisRecorder.submitAllRecords;
+
+    [submitTask continueWithBlock:^id(AWSTask *task) {
+        XCTAssertNotNil(task.error, @"Task should have an error due to invalid pool id.");
+        [expectation fulfill];
+        return nil;
+    }];
+
+    [self waitForExpectationsWithTimeout:5 handler:^(NSError * _Nullable error) {
+        XCTAssertNil(error);
+    }];
+}
+
 - (AWSTask *)getRecords:(NSMutableArray *)returnedRecords shardIterator:(NSString *)shardIterator counter:(int32_t)counter {
     AWSKinesis *kinesis = [AWSKinesis defaultKinesis];
     AWSKinesisGetRecordsInput *getRecordsInput = [AWSKinesisGetRecordsInput new];