Fixed unit test error due to matplotlib