From 8a5e16bfcebe24081147c19d6c4f56c610e12c8b Mon Sep 17 00:00:00 2001
From: Huy Nguyen <huy@pinterest.com>
Date: Fri, 28 Jun 2019 12:05:56 -0700
Subject: [PATCH] Enable experiment in tests

---
 Tests/ASCollectionModernDataSourceTests.mm |  5 +++++
 Tests/ASCollectionViewTests.mm             | 12 +++++++++++-
 Tests/ASCollectionViewThrashTests.mm       | 13 ++++++++++++-
 Tests/ASConfigurationTests.mm              |  3 ++-
 Tests/ASTableViewTests.mm                  | 11 ++++++++++-
 Tests/ASTableViewThrashTests.mm            | 12 +++++++++++-
 6 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/Tests/ASCollectionModernDataSourceTests.mm b/Tests/ASCollectionModernDataSourceTests.mm
index 5d0a7789..e987914f 100644
--- a/Tests/ASCollectionModernDataSourceTests.mm
+++ b/Tests/ASCollectionModernDataSourceTests.mm
@@ -34,6 +34,11 @@
 
 - (void)setUp {
   [super setUp];
+
+  ASConfiguration *config = [ASConfiguration new];
+  config.experimentalFeatures = ASExperimentalOptimizeDataControllerPipeline;
+  [ASConfigurationManager test_resetWithConfiguration:config];
+
   // Default is 2 sections: 2 items in first, 1 item in second.
   sections = [NSMutableArray array];
   [sections addObject:[ASTestSection new]];
diff --git a/Tests/ASCollectionViewTests.mm b/Tests/ASCollectionViewTests.mm
index 3ff5af44..7359aadd 100644
--- a/Tests/ASCollectionViewTests.mm
+++ b/Tests/ASCollectionViewTests.mm
@@ -16,7 +16,9 @@
 #import <OCMock/OCMock.h>
 #import <AsyncDisplayKit/ASCollectionView+Undeprecated.h>
 #import <AsyncDisplayKit/ASDisplayNode+FrameworkPrivate.h>
+
 #import "ASDisplayNodeTestsHelper.h"
+#import "ASTestCase.h"
 
 @interface ASTextCellNodeWithSetSelectedCounter : ASTextCellNode
 
@@ -166,12 +168,20 @@
 
 @end
 
-@interface ASCollectionViewTests : XCTestCase
+@interface ASCollectionViewTests : ASTestCase
 
 @end
 
 @implementation ASCollectionViewTests
 
+- (void)setUp
+{
+  [super setUp];
+  ASConfiguration *config = [ASConfiguration new];
+  config.experimentalFeatures = ASExperimentalOptimizeDataControllerPipeline;
+  [ASConfigurationManager test_resetWithConfiguration:config];
+}
+
 - (void)tearDown
 {
   // We can't prevent the system from retaining windows, but we can at least clear them out to avoid
diff --git a/Tests/ASCollectionViewThrashTests.mm b/Tests/ASCollectionViewThrashTests.mm
index 4650639d..691406f1 100644
--- a/Tests/ASCollectionViewThrashTests.mm
+++ b/Tests/ASCollectionViewThrashTests.mm
@@ -12,9 +12,10 @@
 #import <AsyncDisplayKit/ASCollectionView.h>
 #import <stdatomic.h>
 
+#import "ASTestCase.h"
 #import "ASThrashUtility.h"
 
-@interface ASCollectionViewThrashTests : XCTestCase
+@interface ASCollectionViewThrashTests : ASTestCase
 
 @end
 
@@ -25,8 +26,18 @@
   BOOL _failed;
 }
 
+- (void)setUp
+{
+  [super setUp];
+  ASConfiguration *config = [ASConfiguration new];
+  config.experimentalFeatures = ASExperimentalOptimizeDataControllerPipeline;
+  [ASConfigurationManager test_resetWithConfiguration:config];
+}
+
 - (void)tearDown
 {
+  [super tearDown];
+  
   if (_failed && _update != nil) {
     NSLog(@"Failed update %@: %@", _update, _update.logFriendlyBase64Representation);
   }
diff --git a/Tests/ASConfigurationTests.mm b/Tests/ASConfigurationTests.mm
index c50fa318..c5e533e2 100644
--- a/Tests/ASConfigurationTests.mm
+++ b/Tests/ASConfigurationTests.mm
@@ -59,7 +59,8 @@ static ASExperimentalFeatures features[] = {
   ];
 }
 
-- (ASExperimentalFeatures)allFeatures {
+- (ASExperimentalFeatures)allFeatures
+{
   ASExperimentalFeatures allFeatures = 0;
   for (int i = 0; i < sizeof(features)/sizeof(ASExperimentalFeatures); i++) {
     allFeatures |= features[i];
diff --git a/Tests/ASTableViewTests.mm b/Tests/ASTableViewTests.mm
index e367b0b7..f5e98c27 100644
--- a/Tests/ASTableViewTests.mm
+++ b/Tests/ASTableViewTests.mm
@@ -22,6 +22,7 @@
 #import <AsyncDisplayKit/ASTableView+Undeprecated.h>
 #import <AsyncDisplayKit/ASInternalHelpers.h>
 
+#import "ASTestCase.h"
 #import "ASXCTExtensions.h"
 
 #define NumberOfSections 10
@@ -217,12 +218,20 @@
 
 @end
 
-@interface ASTableViewTests : XCTestCase
+@interface ASTableViewTests : ASTestCase
 @property (nonatomic, retain) ASTableView *testTableView;
 @end
 
 @implementation ASTableViewTests
 
+- (void)setUp
+{
+  [super setUp];
+  ASConfiguration *config = [ASConfiguration new];
+  config.experimentalFeatures = ASExperimentalOptimizeDataControllerPipeline;
+  [ASConfigurationManager test_resetWithConfiguration:config];
+}
+
 - (void)testDataSourceImplementsNecessaryMethods
 {
   ASTestTableView *tableView = [[ASTestTableView alloc] __initWithFrame:CGRectMake(0, 0, 100, 400)
diff --git a/Tests/ASTableViewThrashTests.mm b/Tests/ASTableViewThrashTests.mm
index 8d7064cd..695bf67f 100644
--- a/Tests/ASTableViewThrashTests.mm
+++ b/Tests/ASTableViewThrashTests.mm
@@ -13,9 +13,10 @@
 #import <AsyncDisplayKit/ASTableView+Undeprecated.h>
 #import <stdatomic.h>
 
+#import "ASTestCase.h"
 #import "ASThrashUtility.h"
 
-@interface ASTableViewThrashTests: XCTestCase
+@interface ASTableViewThrashTests: ASTestCase
 @end
 
 @implementation ASTableViewThrashTests
@@ -27,8 +28,17 @@
 
 #pragma mark Overrides
 
+- (void)setUp
+{
+  [super setUp];
+  ASConfiguration *config = [ASConfiguration new];
+  config.experimentalFeatures = ASExperimentalOptimizeDataControllerPipeline;
+  [ASConfigurationManager test_resetWithConfiguration:config];
+}
+
 - (void)tearDown
 {
+  [super tearDown];
   if (_failed && _update != nil) {
     NSLog(@"Failed update %@: %@", _update, _update.logFriendlyBase64Representation);
   }