Fun with XCTestExpectations

XCTestCase+FTTObjectDestructionWaiter.h

#import <XCTest/XCTest.h>

NS_ASSUME_NONNULL_BEGIN

@interface XCTestCase (FTTObjectDestructionWaiter)

- (void)waitForDestructionOfObject:(id)anObject;
- (void)waitForDestructionOfObject:(id)anObject timeout:(NSTimeInterval)timeout;

@end

NS_ASSUME_NONNULL_END

XCTestCase+FTTObjectDestructionWaiter.m

#import "XCTestCase+FTTObjectDestructionWaiter.h"

#import <objc/runtime.h>


static void const *kRunOnDestructionKey = &kRunOnDestructionKey;
static NSTimeInterval const kDefaultTimeout = 2.0;

@interface FTTBlockOnDeallocExecutor: NSObject

@property (nonatomic, readonly) void (^ _Nonnull block)(void);

- (instancetype)initWithBlock:(void (^ _Nonnull)(void))block;

@end


@implementation XCTestCase (FTTObjectDestructionWaiter)

- (void)waitForDestructionOfObject:(id)anObject
{
    [self waitForDestructionOfObject:anObject timeout:kDefaultTimeout];
}

- (void)waitForDestructionOfObject:(id)anObject timeout:(NSTimeInterval)timeout
{
    XCTestExpectation *expectation =
    [self expectationWithDescription:@"waiting for object to be deallocated"];
    
    FTTBlockOnDeallocExecutor *executor = [[FTTBlockOnDeallocExecutor alloc] initWithBlock:^{
        [expectation fulfill];
    }];
    
    objc_setAssociatedObject(anObject,
                             kRunOnDestructionKey,
                             executor,
                             OBJC_ASSOCIATION_RETAIN);
    
    [self waitForExpectations:@[expectation] timeout:timeout];
}

@end


@implementation FTTBlockOnDeallocExecutor

- (instancetype)initWithBlock:(void (^)(void))block
{
    self = [super init];
    
    if (self) {
        _block = [block copy];
    }
    
    return self;
}

- (void)dealloc
{
    _block();
}

@end