EndzeitBegins commented on code in PR #9108:
URL: https://github.com/apache/nifi/pull/9108#discussion_r1694279813


##########
nifi-extension-bundles/nifi-extension-utils/nifi-record-utils/nifi-mock-record-utils/src/main/java/org/apache/nifi/serialization/record/MockSchemaRegistry.java:
##########
@@ -33,40 +33,59 @@
 import java.util.concurrent.ConcurrentMap;
 
 public class MockSchemaRegistry extends AbstractControllerService implements 
SchemaRegistry {
-    private final ConcurrentMap<String, RecordSchema> schemaNameMap = new 
ConcurrentHashMap<>();
+    private final ConcurrentMap<Triple<String, String, Integer>, RecordSchema> 
schemaNameMap = new ConcurrentHashMap<>();
     private final ConcurrentMap<Tuple<Long, Integer>, RecordSchema> 
schemaIdVersionMap = new ConcurrentHashMap<>();
 
     public void addSchema(final String name, final RecordSchema schema) {
-        schemaNameMap.put(name, schema);
+        addSchema(name, null, null, schema);
     }
 
-    RecordSchema retrieveSchemaByName(final SchemaIdentifier schemaIdentifier) 
throws IOException, SchemaNotFoundException {
+    public void addSchema(final String name, final String branch, final 
RecordSchema schema) {
+        addSchema(name, branch, null, schema);
+    }
+
+    public void addSchema(final String name, final Integer version, final 
RecordSchema schema) {
+        addSchema(name, null, version, schema);
+    }
+
+    public void addSchema(final String name, final String branch, final 
Integer version, final RecordSchema schema) {
+        schemaNameMap.put(Triple.of(name, branch, version), schema);
+    }
+
+    RecordSchema retrieveSchemaByName(final SchemaIdentifier schemaIdentifier) 
throws SchemaNotFoundException {
         final Optional<String> schemaName = schemaIdentifier.getName();
-        if (!schemaName.isPresent()) {
-            throw new 
org.apache.nifi.schema.access.SchemaNotFoundException("Cannot retrieve schema 
because Schema Name is not present");
+        if (schemaName.isEmpty()) {
+            throw new 
org.apache.nifi.schema.access.SchemaNotFoundException("Cannot retrieve schema 
because Schema Name is empty");
         }
 
-        return schemaNameMap.get(schemaName.get());
+        if (schemaIdentifier.getBranch().isPresent() && 
schemaIdentifier.getVersion().isPresent()) {
+            return schemaNameMap.get(Triple.of(schemaName.get(), 
schemaIdentifier.getBranch().get(), schemaIdentifier.getVersion().getAsInt()));
+        } else if (schemaIdentifier.getBranch().isPresent() && 
schemaIdentifier.getVersion().isEmpty()) {
+            return schemaNameMap.get(Triple.of(schemaName.get(), 
schemaIdentifier.getBranch().get(), null));
+        } else if (schemaIdentifier.getBranch().isEmpty() && 
schemaIdentifier.getVersion().isPresent()) {
+            return schemaNameMap.get(Triple.of(schemaName.get(), null, 
schemaIdentifier.getVersion().getAsInt()));
+        } else {
+            return schemaNameMap.get(Triple.of(schemaName.get(), null, null));
+        }

Review Comment:
   Thank you for the adjustment. The code is easier to read now from my 
perspective. 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to