diff --git a/tests/Integration/Api/Client/Server/CommandControllerTest.php b/tests/Integration/Api/Client/Server/CommandControllerTest.php
index 8fa106cfc..96cff4ffe 100644
--- a/tests/Integration/Api/Client/Server/CommandControllerTest.php
+++ b/tests/Integration/Api/Client/Server/CommandControllerTest.php
@@ -5,11 +5,11 @@ namespace Pterodactyl\Tests\Integration\Api\Client\Server;
 use Mockery;
 use GuzzleHttp\Psr7\Request;
 use Illuminate\Http\Response;
+use Mockery\MockInterface;
 use Pterodactyl\Models\Server;
 use Pterodactyl\Models\Permission;
 use GuzzleHttp\Exception\BadResponseException;
 use GuzzleHttp\Psr7\Response as GuzzleResponse;
-use Pterodactyl\Repositories\Wings\DaemonCommandRepository;
 use Pterodactyl\Exceptions\Http\Connection\DaemonConnectionException;
 use Pterodactyl\Tests\Integration\Api\Client\ClientApiIntegrationTestCase;
 
@@ -53,12 +53,12 @@ class CommandControllerTest extends ClientApiIntegrationTestCase
     {
         [$user, $server] = $this->generateTestAccount([Permission::ACTION_CONTROL_CONSOLE]);
 
-        $mock = $this->mock(DaemonCommandRepository::class);
-        $mock->expects('setServer')
-            ->with(Mockery::on(fn (Server $value) => $value->is($server)))
-            ->andReturnSelf();
+        $server = \Mockery::mock($server)->makePartial();
+        $server->expects('query->where->firstOrFail')->andReturns($server);
 
-        $mock->expects('send')->with('say Test')->andReturn(new GuzzleResponse());
+        $this->instance(Server::class, $server);
+
+        $server->expects('send')->with('say Test')->andReturn(new GuzzleResponse());
 
         $response = $this->actingAs($user)->postJson("/api/client/servers/$server->uuid/command", [
             'command' => 'say Test',
@@ -75,13 +75,16 @@ class CommandControllerTest extends ClientApiIntegrationTestCase
     {
         [$user, $server] = $this->generateTestAccount();
 
-        $mock = $this->mock(DaemonCommandRepository::class);
-        $mock->expects('setServer->send')->andThrows(
+        $server = \Mockery::mock($server)->makePartial();
+        $server->expects('query->where->firstOrFail')->andReturns($server);
+        $server->expects('send')->andThrows(
             new DaemonConnectionException(
                 new BadResponseException('', new Request('GET', 'test'), new GuzzleResponse(Response::HTTP_BAD_GATEWAY))
             )
         );
 
+        $this->instance(Server::class, $server);
+
         $response = $this->actingAs($user)->postJson("/api/client/servers/$server->uuid/command", [
             'command' => 'say Test',
         ]);