Moved class from test to main
Fixed Lists import