Added simple test. Refactor for testing