add test
make code better
test data