Update test.py

Added element styling code