@Test
public void shouldWriteSuccessAndFailureTagsCorrectly() {

  FeatureSetSpec fs1 =
      FeatureSetSpec.newBuilder()
          .setName("feature_set")
          .setProject("myproject")
          .addEntities(
              EntitySpec.newBuilder()
                  .setName("entity_id_primary")
                  .setValueType(Enum.INT32)
                  .build())
          .addEntities(
              EntitySpec.newBuilder()
                  .setName("entity_id_secondary")
                  .setValueType(Enum.STRING)
                  .build())
          .addFeatures(
              FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build())
          .addFeatures(
              FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build())
          .build();
  FeatureSetSpec fs2 =
      FeatureSetSpec.newBuilder()
          .setName("feature_set_2")
          .setProject("myproject")
          .addEntities(
              EntitySpec.newBuilder()
                  .setName("entity_id_primary")
                  .setValueType(Enum.INT32)
                  .build())
          .addEntities(
              EntitySpec.newBuilder()
                  .setName("entity_id_secondary")
                  .setValueType(Enum.STRING)
                  .build())
          .addFeatures(
              FeatureSpec.newBuilder().setName("feature_1").setValueType(Enum.STRING).build())
          .addFeatures(
              FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.INT64).build())
          .build();
  Map<String, FeatureSetSpec> featureSetSpecs = new HashMap<>();
  featureSetSpecs.put("myproject/feature_set", fs1);
  featureSetSpecs.put("myproject/feature_set_2", fs2);
  List<FeatureRow> input = new ArrayList<>();
  List<FeatureRow> expected = new ArrayList<>();
  for (FeatureSetSpec featureSetSpec : featureSetSpecs.values()) {
    FeatureRow randomRow = TestUtil.createRandomFeatureRow(featureSetSpec);
    input.add(randomRow);
    expected.add(randomRow);
  }
  FeatureRow invalidRow =
      FeatureRow.newBuilder()
          .setFeatureSet(getFeatureSetStringRef(fs1))
          .addFields(
              Field.newBuilder()
                  .setName("feature_1")
                  .setValue(Value.newBuilder().setBoolVal(false).build())
                  .build())
          .build();
  input.add(invalidRow);
  PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView =
      p.apply("StaticSpecs", Create.of(featureSetSpecs)).apply(View.asMultimap());
  PCollectionTuple output =
      p.apply(Create.of(input))
          .setCoder(ProtoCoder.of(FeatureRow.class))
          .apply(
              ProcessAndValidateFeatureRows.newBuilder()
                  .setDefaultProject("myproject")
                  .setFailureTag(FAILURE_TAG)
                  .setSuccessTag(SUCCESS_TAG)
                  .setFeatureSetSpecs(specsView)
                  .build());
  PAssert.that(output.get(SUCCESS_TAG)).containsInAnyOrder(expected);
  PAssert.that(output.get(FAILURE_TAG).apply(Count.globally())).containsInAnyOrder(1L);
  p.run();
}