Fix test cases
New graph class and test