Added some class test to main.py