Added mocks for tests + tests rewritten + some optimizations