Fix ImportError.
More robust error handling + tests.