Revert "Cleanup some custom code for testing"

This reverts commit bc1a34247d65d6b82932d154f6da87456dce97dc.