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


class CreateTrailRequest(Obfuscator):
    target_class = trail_service_pb2.CreateTrailRequest

    def data(self) -> dict:
        request: trail_service_pb2.CreateTrailRequest = self.original

        return {
            'compute_instance_id': request.compute_instance_id,
        }


class CreateTrailResponse(Obfuscator):
    target_class = trail_service_pb2.CreateTrailResponse

    def data(self) -> dict:
        response: trail_service_pb2.CreateTrailResponse = self.original

        return {
            'code': response.code,
            'trail_id': response.trail_id
        }
