Bug fix

Tests working again now.