Orchestrate test case for function that's using threading.Timer

Question

I want to create a test for a function that:

1. Calls an api to create an order
2. Uses a timer to wait for the order to be filled

# This is module 

class Order:
    def __init__(self):
        self._client = Client()

    def open(self, volume: float, type: str):
        try:
            order = self._client.create_order(
                type=type,
                quantity=volume
            )
    
        except Exception as e:
            logger.error(
                f'Exception on open_trade in {symbol} {TradeSide.BUY} \n {e} \n')
            return None
    
    
        t = Timer(1.0, lambda: self._check_order_status(order['orderId']))
        t.start()
         
        return order

    def _check_order_status(self, order_id: str) -> Dict:
        try:
            order = self._client.get_order(orderId=order_id)
          
        except Exception as e:
           
            logger.error(
                f'Exception on getting order status {order_id} {e} ')
            order = None
    
        if order and order['status'] == FILLED:
            
            self.state.update_order(
                self, order)
    
        else:
            t = Timer(2.0, lambda: self._check_order_status(order_id))
            t.start()

To achieve this I mocked the two _client functions:

def test_open(mocker: MockFixture,
                           new_open: Dict, get_open_order: Dict):

    # Arange
    def mock_create_order(self, type, quantity):
        return new_open

    mocker.patch(
        'module.Client.create_order',
        mock_create_order
    )

    def mock_get_order(self, orderId):
        return get_open_order

    mocker.patch(
        'module.Client.get_order',
        mock_get_order
    )

    # Act
    order = Order()
    order.open('BUY', 123)


    # Assert
    assert len(mock_state.open) == 1

The problem is that after starting the Timer, that thread doesn't have the mocked context and it calls the actual class... Any ideas how can I trick the Timer into calling the correct mocked get_order function?

Answer 1

that thread doesn't have the mocked context

That's because test_open has ended and the patched methods have been restored.

Any ideas how can I trick the Timer into calling the correct mocked get_order function?

Here are 3 options. I'd go with option 2.

1. Sleep

This is straightforward, but requires hardcoding the correct time to sleep.

order.open('BUY', 123)
sleep(1.0)  # Add this

2. Track and join the timer thread

This is the clearest way.

t = Timer(1.0, lambda: self._check_order_status(order['orderId']))
t.start()
self._t = t  # Add this

order.open('BUY', 123)
order._t.join()  # Add this

Some may feel that the above leaks a test requirement into the module.
You may want to track the timer thread in a patched Timer class instead:

class TrackedTimer(Timer):
    instance_by_caller_id = {}

    def __init__(self, interval, function, args=None, kwargs=None):
        super().__init__(interval, function, args=args, kwargs=kwargs)

        caller = inspect.currentframe().f_back.f_locals.get('self')
        if caller:
            self.instance_by_caller_id[id(caller)] = self

    @classmethod
    def join_for_caller(cls, caller):
        instance = cls.instance_by_caller_id.get(id(caller))
        if instance:
            instance.join()

mocker.patch(              # Add this
    'polls.module.Timer',  #
    TrackedTimer,          #
)                          #

# Act
order = Order()
order.open('BUY', 123)
TrackedTimer.join_for_caller(order)  # Add this

3. Replace the method on class or instance directly

This allows the execution of the mock method to occur after test_open ends.

Replace the method on the class directly:

# mocker.patch(                         # Replace this
#     'polls.module.Client.get_order',  #
#     mock_get_order                    #
# )                                     #

# Act
Client.get_order = mock_get_order       # with this
order = Order()
order.open('BUY', 123)

Or replace the method on the instance directly:

order = Order()
order._client.get_order = lambda orderId: mock_get_order(order._client, orderId)  # Add this
# setattr(order._client, 'get_order', mock_get_order.__get__(order._client))      # or a bound method
order.open('BUY', 123)