from load.projects.cloud.loadtesting.server.obfuscator.base import Obfuscator
from yandex.cloud.priv.loadtesting.agent.v1 import agent_registration_service_pb2


class RegisterRequest(Obfuscator):
    target_class = agent_registration_service_pb2.RegisterRequest

    def data(self) -> dict:
        request: agent_registration_service_pb2.RegisterRequest = self.original

        return {
            'compute_instance_id': request.compute_instance_id,
        }


class RegisterResponse(Obfuscator):
    target_class = agent_registration_service_pb2.RegisterResponse

    def data(self) -> dict:
        request: agent_registration_service_pb2.RegisterResponse = self.original

        return {
            'agent_instance_id': request.agent_instance_id,
        }


class ExternalAgentRegisterRequest(Obfuscator):
    target_class = agent_registration_service_pb2.ExternalAgentRegisterRequest

    def data(self) -> dict:
        request: agent_registration_service_pb2.ExternalAgentRegisterRequest = self.original

        return {
            'folder_id': request.folder_id,
            'name': request.name,
            'compute_instance_id': request.compute_instance_id,
            'agent_version': request.agent_version,
        }
