Feature/test improvements