fix test

apache · Mar 30, 2022 · 192037b · 192037b
1 parent d8ef1a5
commit 192037b
Showing 1 changed file with 25 additions and 1 deletion.
diff --git a/.../src/main/java/org/apache/pulsar/tests/integration/functions/RemoveAvroFieldFunction.java b/.../src/main/java/org/apache/pulsar/tests/integration/functions/RemoveAvroFieldFunction.java
@@ -56,7 +56,7 @@ public Void process(GenericObject genericObject, Context context) throws Excepti
 
         Schema outputSchema = schema;
         Object outputObject = genericObject.getNativeObject();
-
+        boolean someThingDone = false;
         if (schema instanceof KeyValueSchema && nativeObject instanceof KeyValue)  {
             KeyValueSchema kvSchema = (KeyValueSchema) schema;
 
@@ -96,6 +96,7 @@ public Void process(GenericObject genericObject, Context context) throws Excepti
                     Schema newValueSchema = Schema.NATIVE_AVRO(modified);
                     outputSchema = Schema.KeyValue(keySchema, newValueSchema, kvSchema.getKeyValueEncodingType());
                     outputObject = new KeyValue(originalObject.getKey(), newValue);
+                    someThingDone = true;
                 }
             }
         } else if (schema.getSchemaInfo().getType() == SchemaType.AVRO) {
@@ -125,6 +126,29 @@ public Void process(GenericObject genericObject, Context context) throws Excepti
                 Schema newValueSchema = Schema.NATIVE_AVRO(modified);
                 outputSchema = newValueSchema;
                 outputObject = oo.toByteArray();
+                someThingDone = true;
+            }
+        }
+
+        if (!someThingDone) {
+            // do some processing...
+            final boolean isStruct;
+            switch (currentRecord.getSchema().getSchemaInfo().getType()) {
+                case AVRO:
+                case JSON:
+                case PROTOBUF_NATIVE:
+                    isStruct = true;
+                    break;
+                default:
+                    isStruct = false;
+                    break;
+            }
+            if (isStruct) {
+                // GenericRecord must stay wrapped
+                outputObject = currentRecord.getValue();
+            } else {
+                // primitives and KeyValue must be unwrapped
+                outputObject = nativeObject;
             }
         }
         log.info("output {} schema {}", outputObject, outputSchema);