Polished some code. Defined relevant unit tests.