Add Test

Made some modification to add the jUnit test