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


class SayHelloRequest(Obfuscator):
    target_class = greeter_service_pb2.SayHelloRequest

    def data(self) -> dict:
        request: greeter_service_pb2.SayHelloRequest = self.original

        return {
            'tank_id': request.tank_id,
        }


class SayHelloResponse(Obfuscator):
    target_class = greeter_service_pb2.SayHelloResponse

    def data(self) -> dict:
        response: greeter_service_pb2.SayHelloResponse = self.original

        return {
            'code': response.code,
        }
