Fixing bugs, wrote a test