Add import for TestProbe

The example for TestProbe doesn't seem to work very well without the TestProbe class imported. :-(