Fix step definition tests for presence of heading and minor refactoring.

@andrewgarner - can you review the fix to the step definition just to check I've done the right thing.