Cleaned up test code (using API)