Search code examples
unit-testingruststubborrow-checker

How to inject a stub for a trait object into a type in Rust and keep a reference to it?


I'm trying to develop some kind of time tracking CLI tool in Rust. I try to cover most of the code with unit tests and got stuck about how to use stub objects with Rusts ownership restrictions.

I have a type Tracker which has functions to start/stop tracking time. These functions query the current system time.

To make this design testable I introduced the trait TimeService which provides the current timestamp. I have the "real" implementation SystemTimeService which returns the current system time and a fake implementation FakeTimeService for tests where the time can be set from the outside.

This is the current implementation:

pub struct Tracker {
    // ...
    time_service: Box<TimeService>,
}

trait TimeService {
    fn time(&self) -> u64;
}

Here is the test I'd like to implement:

#[test]
fn tracker_end_when_called_tracks_total_time() {
    let (tracker, time_service) = init_tracker();

    time_service.set_seconds(0);
    tracker.start("feature1");
    time_service.set_seconds(10);
    tracker.end();

    assert_eq!(tracker.total_time, 10);
}

fn init_tracker() -> (Tracker, &FakeTimeService) {
    let time_service = Box::new(FakeTimeService::new());
    let tracker = Tracker::with_time_service(time_service);

    // does not compile, as time service's ownership has been
    // moved to the tracker.
    (tracker, &*time_service)
}

The problem is that I don't know how to access the fake time service from inside the unit test as it's ownership is taken by the tracker.

I can imagine multiple solutions:

  1. Use Rc instead of Box inside the tracker for shared ownership of the time service
  2. Make the Tracker generic and add a type argument for the used TimeTracker implementation
  3. Add lifetimes to the init_tracker function

I do not like using solution 1 as this would not express the idea that the service is part of the tracker (seems to violate encapsulation).

Solution 2 is probably viable but in that case I'd need to make the TimeService trait and the used implementations public which I'd also like to avoid.

So the most promising solution without changing the design seems to be lifetimes. Is it possible to add lifetimes to the variables in such a way that the init_tracker function can return the tracker and the fake time service?

Are there any other solutions/best practices?


Solution

  • You can make the FakeTimeService type cloneable, making its state shared across multiple cloned instances:

    use std::rc::Rc;
    
    #[derive(Clone)]
    struct FakeTimeService {
        state: Rc<FakeTimeServiceState>,
    }
    
    impl TimeService for FakeTimeService { ... }
    
    impl FakeTimeService {
        fn state(&self) -> &FakeTimeServiceState {
            *self.state
        }
    }
    
    fn init_tracker() -> (Tracker, FakeTimeService) {
        let time_service = FakeTimeService::new();
        let tracker = Tracker::with_time_service(Box::new(time_service.clone()));
        (tracker, time_service)
    }
    

    Now you'll be able to independently modify the fake time service state, while keeping the interface and implementation of the Tracker the same as before.

    Most likely you'll want to have some mutable state inside FakeTimeService in order to set up your test preconditions. Therefore, you will probably need to use some kind of internal mutability for the FakeTimeService state:

    use std::rc::Rc;
    use std::cell::{RefCell, RefMut};
    
    #[derive(Clone)]
    struct FakeTimeService {
        state: Rc<RefCell<FakeTimeServiceState>>,
    }
    
    impl FakeTimeService {
        fn state_mut(&self) -> RefMut<FakeTimeServiceState> {
            self.state.borrow_mut()
        }
    }
    
    struct FakeTimeServiceState {
        init_value: i32,
    }
    
    impl FakeTimeServiceState {
        fn set_initial_value(&mut self, x: i32) {
            self.init_value = x;
        }
    }
    
    let (tracker, time_service) = init_tracker();
    time_service.state_mut().set_initial_value(123);