Enhance testCustomFallbackHosts

ably · Oct 11, 2016 · f88d683 · f88d683
1 parent ab33a48
commit f88d683
Showing 1 changed file with 31 additions and 39 deletions.
diff --git a/Tests/ARTFallbackTest.m b/Tests/ARTFallbackTest.m
@@ -44,46 +44,38 @@ - (void)testAllHostsIncludedOnce {
 }
 
 - (void)testCustomFallbackHosts {
-    NSArray *customHosts = @[@"test.ably.com",
-                             @"test2.ably.com",
-                             @"test3.ably.com",
-                             @"test4.ably.com",
-                             @"test5.ably.com",
-                             @"test6.ably.com",
-                             @"test7.ably.com",
-                             @"test8.ably.com",
-                             @"test9.ably.com",
-                             @"test10.ably.com",
-                             @"test11.ably.com",
-                             @"test12.ably.com",
-                             @"test13.ably.com",
-                             @"test14.ably.com"];
-    ARTFallback *f = [[ARTFallback alloc] initWithFallbackHosts:customHosts];
-
-    NSSet *customSet = [NSSet setWithArray:customHosts];
-    NSMutableArray *hostsRandomised = [NSMutableArray array];
-    for(int i=0;i < [customHosts count]; i++) {
-        [hostsRandomised addObject:[f popFallbackHost]];
-    }
-
-    //popping after all hosts are exhausted returns nil
-    XCTAssertTrue([f popFallbackHost] == nil);
-
-    // all fallback hosts are used in artfallback
-    XCTAssertEqual([hostsRandomised count], [customHosts count]);
-    bool inOrder = true;
-    for(int i=0;i < [customHosts count]; i++) {
-        if(![[customHosts objectAtIndex:i] isEqualToString:[hostsRandomised objectAtIndex:i]]) {
-            inOrder = false;
-            break;
-        }
+    __weak int (^originalARTFallback_getRandomHostIndex)(int) = ARTFallback_getRandomHostIndex;
+    @try {
+        ARTFallback_getRandomHostIndex = ^() {
+            __block NSArray *hostIndexes = @[@1, @2, @0, @1, @0, @0];
+            __block int i = 0;
+            return ^int(int count) {
+                NSNumber *hostIndex = hostIndexes[i];
+                i++;
+                return hostIndex.intValue;
+            };
+        }();
+
+        NSArray *customHosts = @[@"testA.ably.com",
+                                 @"testB.ably.com",
+                                 @"testC.ably.com",
+                                 @"testD.ably.com",
+                                 @"testE.ably.com",
+                                 @"testF.ably.com"];
+
+        ARTFallback *f = [[ARTFallback alloc] initWithFallbackHosts:customHosts];
+
+        XCTAssertEqualObjects([f popFallbackHost], @"testF.ably.com");
+        XCTAssertEqualObjects([f popFallbackHost], @"testC.ably.com");
+        XCTAssertEqualObjects([f popFallbackHost], @"testE.ably.com");
+        XCTAssertEqualObjects([f popFallbackHost], @"testA.ably.com");
+        XCTAssertEqualObjects([f popFallbackHost], @"testD.ably.com");
+        XCTAssertEqualObjects([f popFallbackHost], @"testB.ably.com");
+
+        XCTAssertEqual([f popFallbackHost], nil);
     }
-    //check artfallback randomises the order.
-    XCTAssertFalse(inOrder);
-
-    //every member of fallbacks hosts are in the list of custom hosts
-    for(int i=0;i < [hostsRandomised count]; i++) {
-        XCTAssertTrue([customSet containsObject:[hostsRandomised objectAtIndex:i]]);
+    @finally {
+        ARTFallback_getRandomHostIndex = originalARTFallback_getRandomHostIndex;
     }
 }