diff --git a/tests/src/Unit/Commands/PubSubCommandsTest.php b/tests/src/Unit/Commands/PubSubCommandsTest.php
index 03ccb34..697a042 100644
--- a/tests/src/Unit/Commands/PubSubCommandsTest.php
+++ b/tests/src/Unit/Commands/PubSubCommandsTest.php
@@ -8,11 +8,11 @@
 use Drupal\helfi_api_base\Azure\PubSub\PubSubManagerInterface;
 use Drupal\helfi_api_base\Drush\Commands\PubSubCommands;
 use Drush\Commands\DrushCommands;
+use Drush\Style\DrushStyle;
 use Prophecy\Argument;
 use Prophecy\PhpUnit\ProphecyTrait;
 use Symfony\Component\Console\Input\InputInterface;
 use Symfony\Component\Console\Output\OutputInterface;
-use Symfony\Component\Console\Style\SymfonyStyle;
 use WebSocket\ConnectionException;
 use WebSocket\TimeoutException;
 
@@ -31,7 +31,7 @@ public function testListen() : void {
     $expectedMessage = '{"message":"test"}';
     $output = $this->prophesize(OutputInterface::class);
     $input = $this->prophesize(InputInterface::class);
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->writeln(Argument::containingString('Received message'))
       ->shouldBeCalledTimes(PubSubCommands::MAX_MESSAGES);
     $io->writeln(Argument::containingString('Received maximum number of messages'))
@@ -51,7 +51,7 @@ public function testListen() : void {
   public function testExceptionOutput() : void {
     $output = $this->prophesize(OutputInterface::class);
     $input = $this->prophesize(InputInterface::class);
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->writeln('Invalid json: Syntax error')->shouldBeCalledTimes(PubSubCommands::MAX_MESSAGES);
     $io->writeln(Argument::containingString('Received maximum number of messages'))
       ->shouldBeCalledTimes(1);
@@ -70,7 +70,7 @@ public function testExceptionOutput() : void {
   public function testTimeoutException() : void {
     $output = $this->prophesize(OutputInterface::class);
     $input = $this->prophesize(InputInterface::class);
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->writeln(Argument::containingString('Received maximum number of messages'))
       ->shouldBeCalledTimes(1);
 
diff --git a/tests/src/Unit/Commands/RevisionCommandsTest.php b/tests/src/Unit/Commands/RevisionCommandsTest.php
index 335fa60..ad7a133 100644
--- a/tests/src/Unit/Commands/RevisionCommandsTest.php
+++ b/tests/src/Unit/Commands/RevisionCommandsTest.php
@@ -13,12 +13,12 @@
 use Drupal\helfi_api_base\Commands\RevisionCommands;
 use Drupal\helfi_api_base\Entity\Revision\RevisionManager;
 use Drush\Commands\DrushCommands;
+use Drush\Style\DrushStyle;
 use Prophecy\Argument;
 use Prophecy\PhpUnit\ProphecyTrait;
 use Prophecy\Prophecy\ObjectProphecy;
 use Symfony\Component\Console\Input\InputInterface;
 use Symfony\Component\Console\Output\OutputInterface;
-use Symfony\Component\Console\Style\SymfonyStyle;
 
 /**
  * Tests revision commands.
@@ -114,7 +114,7 @@ private function getConnectionMock(array $expected, ?int $id = NULL) : Connectio
    * Test command with invalid entity type.
    */
   public function testInvalidEntityType() : void {
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->writeln(Argument::containingString('Given entity type is not supported.'))
       ->shouldBeCalled();
 
@@ -130,7 +130,7 @@ public function testInvalidEntityType() : void {
    * Test delete without any entities.
    */
   public function testNoEntities() : void {
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->isDecorated()
       ->willReturn(TRUE);
     $io->getVerbosity()
@@ -148,7 +148,7 @@ public function testNoEntities() : void {
    * Tests delete method with proper data.
    */
   public function testDelete() : void {
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->isDecorated()
       ->willReturn(TRUE);
     $io->getVerbosity()
@@ -170,7 +170,7 @@ public function testDelete() : void {
    * Tests delete method with optional ID.
    */
   public function testDeleteWithId() : void {
-    $io = $this->prophesize(SymfonyStyle::class);
+    $io = $this->prophesize(DrushStyle::class);
     $io->isDecorated()
       ->willReturn(TRUE);
     $io->getVerbosity()