Revert "Revert "Fixing test""

This reverts commit acc7d061f9f1faef375947b77fe902241a8581c4.