Improve test.py
Merge pull request #1 from Elzor/improve_tests

Improve test.py