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


class ClaimAgentStatusRequest(Obfuscator):
    target_class = agent_service_pb2.ClaimAgentStatusRequest

    def data(self) -> dict:
        request: agent_service_pb2.ClaimAgentStatusRequest = self.original

        return {
            'agent_instance_id': request.agent_instance_id,
            'status': agent_service_pb2.ClaimAgentStatusRequest.Status.Name(request.status),
        }


class ClaimAgentStatusResponse(Obfuscator):
    target_class = agent_service_pb2.ClaimAgentStatusResponse

    def data(self) -> dict:
        response: agent_service_pb2.ClaimAgentStatusResponse = self.original

        return {
            'code': response.code,
        }
