Update tests.py

This works
Merge pull request #4 from srujank1/patch-4

Update tests.py