Update Tests for main.py

Tests are outdated and need to be updated. If there is a way to write them so that they work with synthetic examples this would be ideal.