oto/apps/runner/test/oto_agent_registration_test.dart

import 'dart:async';
import 'dart:convert';
import 'dart:io';

import 'package:test/test.dart';
import 'package:oto/oto/agent/agent_config.dart';
import 'package:oto/oto/agent/agent_runner.dart';
import 'package:oto/oto/agent/edge_registration_client.dart';

class FakeEdgeAgentSession implements EdgeAgentSession {
  @override
  final RegistrationResult result;
  bool closed = false;

  FakeEdgeAgentSession(this.result);

  @override
  Future<void> close() async {
    closed = true;
  }
}

class FakeEdgeRegistrationClient extends RegistrationClient {
  final RegistrationResult Function(AgentConfig config) _onRegister;
  FakeEdgeAgentSession? lastSession;

  FakeEdgeRegistrationClient(this._onRegister);

  @override
  Future<EdgeAgentSession> openSession(AgentConfig config) async {
    final session = FakeEdgeAgentSession(_onRegister(config));
    lastSession = session;
    return session;
  }
}

void main() {
  const validBootstrapYaml = '''
agent:
  id: "agent-123"
  alias: "my-agent"
  enrollment_token: "token-456"
edge:
  url: "127.0.0.1:8080"
runtime:
  install_dir: "/usr/bin"
  workspace_root: "/var/oto"
  log_dir: "/var/log/oto"
''';

  late AgentConfig validConfig;

  setUp(() {
    validConfig = AgentConfig.fromYamlContent(validBootstrapYaml);
  });

  group('EdgeEndpoint', () {
    test('EdgeEndpoint parses host and port from config edge url', () {
      final endpoint1 = EdgeEndpoint.parse('127.0.0.1:8080');
      expect(endpoint1.host, '127.0.0.1');
      expect(endpoint1.port, 8080);

      final endpoint2 = EdgeEndpoint.parse('http://edge-server:9000');
      expect(endpoint2.host, 'edge-server');
      expect(endpoint2.port, 9000);

      final endpoint3 = EdgeEndpoint.parse('https://secure-edge');
      expect(endpoint3.host, 'secure-edge');
      expect(endpoint3.port, 443);

      final endpoint4 = EdgeEndpoint.parse('edge-no-port');
      expect(endpoint4.host, 'edge-no-port');
      expect(endpoint4.port, 80);
    });
  });

  group('EdgeRegistrationClient.register one-shot', () {
    test('register returns session result and closes the session', () async {
      final fakeClient = FakeEdgeRegistrationClient((config) {
        return RegistrationResult.accepted('node-123', 'alias-123', {
          'concurrency': 1,
        });
      });

      final result = await fakeClient.register(validConfig);

      expect(result.accepted, isTrue);
      expect(result.nodeId, 'node-123');
      expect(fakeClient.lastSession, isNotNull);
      expect(fakeClient.lastSession!.closed, isTrue);
    });
  });

  group('OtoServerRegistrationClient', () {
    test('builds OTO Server registration request from config', () {
      final client = OtoServerRegistrationClient(
        commandTypes: ['Shell', 'Git'],
      );

      final request = client.buildRegisterRunnerRequest(validConfig);

      expect(request.enrollmentToken, 'token-456');
      expect(request.runnerId, 'agent-123');
      expect(request.alias, 'my-agent');
      expect(request.protocolVersion, otoRunnerProtocolVersion);
      expect(request.capability.name, otoRunnerCapabilityName);
      expect(request.commandCatalog.commandTypes, ['Shell', 'Git']);
    });

    test('posts registration request to OTO Server endpoint', () async {
      final server = await HttpServer.bind(InternetAddress.loopbackIPv4, 0);
      final captured = Completer<Map<String, dynamic>>();
      final serverSub = server.listen((request) async {
        expect(request.method, 'POST');
        if (request.uri.path == '/api/v1/runners/register') {
          final body = await utf8.decoder.bind(request).join();
          captured.complete(jsonDecode(body) as Map<String, dynamic>);
          request.response
            ..headers.contentType = ContentType.json
            ..write(
              jsonEncode({
                'accepted': true,
                'runner_id': 'agent-123',
                'alias': 'my-agent',
              }),
            );
        } else {
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({'success': true}));
        }
        await request.response.close();
      });

      try {
        final config = AgentConfig(
          agent: validConfig.agent,
          server: ServerConnectionConfig(
            url: 'http://${server.address.host}:${server.port}',
          ),
          runtime: validConfig.runtime,
        );
        final client = OtoServerRegistrationClient(
          commandTypes: ['Shell', 'Git'],
        );

        final result = await client.register(config);

        expect(result.accepted, isTrue);
        expect(result.runnerId, 'agent-123');
        final payload = await captured.future;
        expect(payload['enrollment_token'], 'token-456');
        expect(payload['runner_id'], 'agent-123');
        expect(payload['alias'], 'my-agent');
        expect(payload['protocol_version'], otoRunnerProtocolVersion);
        expect(payload['capability'], containsPair('name', 'oto-runner'));
        expect(
          payload['command_catalog'],
          containsPair('command_types', ['Shell', 'Git']),
        );
      } finally {
        await serverSub.cancel();
        await server.close(force: true);
      }
    });

    test('sends first heartbeat and disconnects on close', () async {
      final server = await HttpServer.bind(InternetAddress.loopbackIPv4, 0);
      final registerReceived = Completer<void>();
      final heartbeatReceived = Completer<void>();
      final disconnectReceived = Completer<void>();

      final serverSub = server.listen((request) async {
        if (request.method == 'POST' && request.uri.path == '/api/v1/runners/register') {
          registerReceived.complete();
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({
              'accepted': true,
              'runner_id': 'agent-123',
              'alias': 'my-agent',
            }));
        } else if (request.method == 'POST' && request.uri.path == '/api/v1/runners/agent-123/heartbeat') {
          heartbeatReceived.complete();
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({
              'success': true,
            }));
        } else if (request.method == 'POST' && request.uri.path == '/api/v1/runners/agent-123/disconnect') {
          disconnectReceived.complete();
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({
              'success': true,
            }));
        }
        await request.response.close();
      });

      try {
        final config = AgentConfig(
          agent: validConfig.agent,
          server: ServerConnectionConfig(
            url: 'http://${server.address.host}:${server.port}',
          ),
          runtime: validConfig.runtime,
        );
        final client = OtoServerRegistrationClient(
          commandTypes: ['Shell', 'Git'],
          heartbeatInterval: const Duration(milliseconds: 100),
        );

        final session = await client.openSession(config);

        // First heartbeat should be sent automatically
        await heartbeatReceived.future.timeout(const Duration(seconds: 2));

        // Close session
        await session.close();

        // Disconnect should be sent
        await disconnectReceived.future.timeout(const Duration(seconds: 2));
      } finally {
        await serverSub.cancel();
        await server.close(force: true);
      }
    });

    test('does not send heartbeat or disconnect on rejected registration', () async {
      final server = await HttpServer.bind(InternetAddress.loopbackIPv4, 0);
      var heartbeatReceived = false;
      var disconnectReceived = false;

      final serverSub = server.listen((request) async {
        if (request.method == 'POST' && request.uri.path == '/api/v1/runners/register') {
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({
              'accepted': false,
              'reject_reason': 'Invalid token',
            }));
        } else if (request.uri.path.contains('heartbeat')) {
          heartbeatReceived = true;
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({'success': true}));
        } else if (request.uri.path.contains('disconnect')) {
          disconnectReceived = true;
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({'success': true}));
        }
        await request.response.close();
      });

      try {
        final config = AgentConfig(
          agent: validConfig.agent,
          server: ServerConnectionConfig(
            url: 'http://${server.address.host}:${server.port}',
          ),
          runtime: validConfig.runtime,
        );
        final client = OtoServerRegistrationClient(
          commandTypes: ['Shell', 'Git'],
          heartbeatInterval: const Duration(milliseconds: 50),
        );

        final session = await client.openSession(config);
        expect(session.result.accepted, isFalse);

        // Wait a short duration to ensure no heartbeat is sent
        await Future<void>.delayed(const Duration(milliseconds: 200));
        expect(heartbeatReceived, isFalse);

        await session.close();
        // Wait a short duration to ensure no disconnect is sent
        await Future<void>.delayed(const Duration(milliseconds: 100));
        expect(disconnectReceived, isFalse);
      } finally {
        await serverSub.cancel();
        await server.close(force: true);
      }
    });

    test('session close cancels future heartbeats', () async {
      final server = await HttpServer.bind(InternetAddress.loopbackIPv4, 0);
      var heartbeatCount = 0;

      final serverSub = server.listen((request) async {
        if (request.method == 'POST' && request.uri.path == '/api/v1/runners/register') {
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({
              'accepted': true,
              'runner_id': 'agent-123',
            }));
        } else if (request.uri.path.contains('heartbeat')) {
          heartbeatCount++;
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({'success': true}));
        } else if (request.uri.path.contains('disconnect')) {
          request.response
            ..headers.contentType = ContentType.json
            ..write(jsonEncode({'success': true}));
        }
        await request.response.close();
      });

      try {
        final config = AgentConfig(
          agent: validConfig.agent,
          server: ServerConnectionConfig(
            url: 'http://${server.address.host}:${server.port}',
          ),
          runtime: validConfig.runtime,
        );
        final client = OtoServerRegistrationClient(
          commandTypes: ['Shell', 'Git'],
          heartbeatInterval: const Duration(milliseconds: 50),
        );

        final session = await client.openSession(config);

        // Wait for first heartbeat(s)
        await Future<void>.delayed(const Duration(milliseconds: 150));
        final countBeforeClose = heartbeatCount;
        expect(countBeforeClose, greaterThan(0));

        await session.close();

        // Wait longer to see if any more heartbeats occur
        await Future<void>.delayed(const Duration(milliseconds: 150));
        expect(heartbeatCount, equals(countBeforeClose));
      } finally {
        await serverSub.cancel();
        await server.close(force: true);
      }
    });
  });

  group('AgentRunner', () {
    test('AgentRunner maps config token to register request', () async {
      AgentConfig? capturedConfig;
      final fakeClient = FakeEdgeRegistrationClient((config) {
        capturedConfig = config;
        return RegistrationResult.accepted('node-123', 'alias-123', {
          'concurrency': 1,
        });
      });

      final runner = DefaultAgentRunner(
        client: fakeClient,
        onLog: (msg) {},
        waitForShutdown: () async {},
      );
      await runner.run(validConfig);

      expect(capturedConfig, isNotNull);
      expect(capturedConfig!.agent.enrollmentToken, 'token-456');
    });

    test('AgentRunner keeps session open until shutdown then closes', () async {
      final fakeClient = FakeEdgeRegistrationClient((config) {
        return RegistrationResult.accepted('node-123', 'alias-123', {
          'concurrency': 1,
        });
      });

      final shutdown = Completer<void>();
      final runner = DefaultAgentRunner(
        client: fakeClient,
        onLog: (msg) {},
        waitForShutdown: () => shutdown.future,
      );

      final runFuture = runner.run(validConfig);
      // 종료 신호 전에는 session이 살아 있어야 한다.
      await Future<void>.delayed(Duration.zero);
      expect(fakeClient.lastSession, isNotNull);
      expect(fakeClient.lastSession!.closed, isFalse);

      shutdown.complete();
      await runFuture;
      expect(fakeClient.lastSession!.closed, isTrue);
    });

    test(
      'AgentRunner reports rejected registration and closes session',
      () async {
        final fakeClient = FakeEdgeRegistrationClient((config) {
          return RegistrationResult.rejected('Invalid token');
        });

        final runner = DefaultAgentRunner(
          client: fakeClient,
          onLog: (msg) {},
          waitForShutdown: () async {},
        );
        await expectLater(
          () => runner.run(validConfig),
          throwsA(
            isA<RegistrationException>().having(
              (e) => e.message,
              'message',
              'Invalid token',
            ),
          ),
        );
        expect(fakeClient.lastSession, isNotNull);
        expect(fakeClient.lastSession!.closed, isTrue);
      },
    );
  });
}