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


class ClaimTankStatusRequest(Obfuscator):
    target_class = tank_service_pb2.ClaimTankStatusRequest

    def data(self) -> dict:
        request: tank_service_pb2.ClaimTankStatusRequest = self.original

        return {
            'compute_instance_id': request.compute_instance_id,
            'status': tank_service_pb2.ClaimTankStatusRequest.Status.Name(request.status),
        }


class ClaimTankStatusResponse(Obfuscator):
    target_class = tank_service_pb2.ClaimTankStatusResponse

    def data(self) -> dict:
        response: tank_service_pb2.ClaimTankStatusResponse = self.original

        return {
            'code': response.code,
        }
