Adding unittests for the new plotting functions + small bug fix and cosmetic changes in these functions.