WP-LIT-MCP/tests/Unit/Admin/ConnectionUsageTest.php

<?php

declare(strict_types=1);

namespace LIT\MCP\Tests\Unit\Admin;

use Brain\Monkey\Functions;
use LIT\MCP\Admin\ConnectionUsage;
use LIT\MCP\Tests\Unit\TestCase;
use WP_Error;
use WP_HTTP_Response;
use WP_REST_Request;
use WP_REST_Response;

require_once __DIR__ . '/Fixtures/HttpResponses.php';

final class ConnectionUsageTest extends TestCase
{
    /** @var list<array{int, string}> */
    private array $recorded = [];

    private function usage(?array $claims): ConnectionUsage
    {
        return new ConnectionUsage(
            static fn (string $token): ?array => $token === 'valid-token' ? $claims : null,
            function (int $userId, string $uuid): void {
                $this->recorded[] = [$userId, $uuid];
            }
        );
    }

    private function request(string $route = '/mcp/mcp-oauth-server', string $auth = 'Bearer valid-token'): WP_REST_Request
    {
        return new WP_REST_Request($route, ['authorization' => $auth]);
    }

    public function testRegistersAfterTheRouteCallback(): void
    {
        $usage = $this->usage([]);
        Functions\expect('add_filter')->once()->with('rest_request_after_callbacks', [$usage, 'record'], 10, 3);

        $usage->register();
    }

    public function testRecordsUsageOfTheTokensApplicationPasswordForTheAuthenticatedUser(): void
    {
        Functions\when('get_current_user_id')->justReturn(7);
        $response = ['jsonrpc' => '2.0'];

        $result = $this->usage(['sub' => 7, 'app_pass_id' => 'uuid-1'])->record($response, [], $this->request());

        self::assertSame($response, $result);
        self::assertSame([[7, 'uuid-1']], $this->recorded);
    }

    public function testIgnoresOtherRoutesFailedRequestsAndForeignOrInvalidTokens(): void
    {
        Functions\when('get_current_user_id')->justReturn(7);
        $claims = ['sub' => 7, 'app_pass_id' => 'uuid-1'];

        $this->usage($claims)->record([], [], $this->request('/wp/v2/posts'));
        $this->usage($claims)->record(new WP_Error('unauthorized', 'x'), [], $this->request());
        $this->usage($claims)->record([], [], $this->request('/mcp/mcp-oauth-server', 'Basic abc'));
        $this->usage($claims)->record([], [], $this->request('/mcp/mcp-oauth-server', 'Bearer forged'));
        $this->usage(['sub' => 8, 'app_pass_id' => 'uuid-1'])->record([], [], $this->request());
        $this->usage(['sub' => 7])->record([], [], $this->request());

        self::assertSame([], $this->recorded);
    }

    /** @dataProvider failedStatuses */
    public function testDoesNotRecordRequestsAnsweredWithAnErrorStatus(int $status): void
    {
        Functions\when('get_current_user_id')->justReturn(7);
        $claims = ['sub' => 7, 'app_pass_id' => 'uuid-1'];
        // mcp-oauth answers auth failures with a WP_REST_Response (OAuthHttpTransport::build_auth_error_response).
        $rest = new WP_REST_Response(['error' => 'invalid_token'], $status);

        self::assertSame($rest, $this->usage($claims)->record($rest, [], $this->request()));
        $this->usage($claims)->record(new WP_HTTP_Response(null, $status), [], $this->request());

        self::assertSame([], $this->recorded);
    }

    /** @return list<list<int>> */
    public static function failedStatuses(): array
    {
        return [[401], [403], [500]];
    }

    public function testRecordsSuccessfulRestResponses(): void
    {
        Functions\when('get_current_user_id')->justReturn(7);
        $claims = ['sub' => 7, 'app_pass_id' => 'uuid-1'];

        $this->usage($claims)->record(new WP_REST_Response([], 200), [], $this->request());
        $this->usage($claims)->record(new WP_REST_Response([], 202), [], $this->request());

        self::assertSame([[7, 'uuid-1'], [7, 'uuid-1']], $this->recorded);
    }

    public function testNeverBreaksTheResponseWhenDecodingOrRecordingFails(): void
    {
        Functions\when('get_current_user_id')->justReturn(7);
        $usage = new ConnectionUsage(
            static fn (): ?array => throw new \RuntimeException('bad secret'),
            static function (): void {
            }
        );

        self::assertSame(['ok'], $usage->record(['ok'], [], $this->request()));
    }
}