Updates to unit tests
Draft of stage4.py and stage4 test