Fixed test a bit more (i.e. cleanup unneeded stuff)
Separated test into separate files
Added a few tests and fixed entities.py accordingly