Fixed tests for upstream changes