package com.google.swarm.tokenization.common;

import static org.junit.Assert.*;

import com.google.privacy.dlp.v2.ContentItem;
import com.google.privacy.dlp.v2.DeidentifyContentRequest;
import com.google.privacy.dlp.v2.Table;
import com.google.privacy.dlp.v2.Value;
import org.apache.beam.sdk.options.ValueProvider;
import org.junit.Test;

public class DLPTokenizationDoFnTest {
  @Test
  public void testSetInspectTemplateExist() {
    DLPTokenizationDoFn dlp =
        new DLPTokenizationDoFn(
            "Project Name",
            ValueProvider.StaticValueProvider.of("DeidentifyTemplateName"),
            ValueProvider.StaticValueProvider.of("IdentifyTemplateName"));

    assertEquals(dlp.getInspectTemplateExist(), false);
    dlp.setInspectTemplateExist();
    assertEquals(dlp.getInspectTemplateExist(), true);
  }

  @Test
  public void testBuildDeidentifyContentRequest() {
    DLPTokenizationDoFn dlp =
        new DLPTokenizationDoFn(
            "Project Name",
            ValueProvider.StaticValueProvider.of("DeidentifyTemplateName"),
            ValueProvider.StaticValueProvider.of("IdentifyTemplateName"));

    ContentItem contentItem = ContentItem.newBuilder().build();

    dlp.setInspectTemplateExist();
    DeidentifyContentRequest request = dlp.buildDeidentifyContentRequest(contentItem);

    assertEquals(request.getParent(), "projects/Project Name");
    assertEquals(request.getDeidentifyTemplateName(), "DeidentifyTemplateName");
    assertEquals(request.getInspectTemplateName(), "IdentifyTemplateName");
  }

  @Test
  public void testConvertTableRowToRow() {
    DLPTokenizationDoFn dlp =
        new DLPTokenizationDoFn(
            "Project Name",
            ValueProvider.StaticValueProvider.of("DeidentifyTemplateName"),
            ValueProvider.StaticValueProvider.of("IdentifyTemplateName"));

    String[] header = {"header0", "header1"};
    String key = "Key name";
    Table.Row.Builder tableRowBuilder = Table.Row.newBuilder();
    tableRowBuilder.addValues(0, Value.newBuilder().setStringValue("value0"));
    tableRowBuilder.addValues(1, Value.newBuilder().setStringValue("value1"));
    Table.Row row = tableRowBuilder.build();

    Row result = dlp.convertTableRowToRow(header, key, row);
    assertEquals(result.getTableId(), key);
    assertEquals(result.getHeader()[0], "header0");
    assertEquals(result.getValue()[1], "value1");
  }
}