Cleaned up code and changed tests.
* imports according to PEP 8
* changed tests to test the correct component (was testing the wrong
component before)