diff --git a/lib/babe/babe_test.go b/lib/babe/babe_test.go
index 03a3fd938c..776ba26817 100644
--- a/lib/babe/babe_test.go
+++ b/lib/babe/babe_test.go
@@ -92,7 +92,7 @@ func createTestService(t *testing.T, cfg *ServiceConfig) *Service {
 	}
 
 	cfg.BlockImportHandler = new(mocks.BlockImportHandler)
-	cfg.BlockImportHandler.(*mocks.BlockImportHandler).On("HandleBlockImport", mock.AnythingOfType("*types.Block"), mock.AnythingOfType("*storage.TrieState")).Return(nil)
+	cfg.BlockImportHandler.(*mocks.BlockImportHandler).On("HandleBlockProduced", mock.AnythingOfType("*types.Block"), mock.AnythingOfType("*storage.TrieState")).Return(nil)
 
 	if cfg.Keypair == nil {
 		cfg.Keypair, err = sr25519.GenerateKeypair()