SK-2646: fix tokenize reponse handling

Author: saileshwar-skyflow

v3/src/main/java/com/skyflow/utils/Utils.java

@@ -445,9 +445,6 @@ public static TokenizeResponse formatTokenizeResponse(
             com.skyflow.generated.rest.resources.flowservice.requests.V1FlowTokenizeRequest batchRequest,
             int batchNumber, int batchSize) {
         if (response != null && response.getResponse().isPresent()) {
-            // The API returns a flat list — one entry per (value x tokenGroupName).
-            // We group them by input record position using the request's tokenGroupNames count
-            // so that records with identical values are still treated as separate entries.
             List<com.skyflow.generated.rest.types.V1FlowTokenizeResponseObject> flatList =
                     response.getResponse().get();
             List<com.skyflow.generated.rest.types.V1FlowTokenizeRequestObject> requestData =
@@ -456,20 +453,19 @@ public static TokenizeResponse formatTokenizeResponse(
             List<TokenizeSuccess> successRecords = new ArrayList<>();
             List<ErrorRecord> errorRecords = new ArrayList<>();
 
+
             int flatIndex = 0;
             for (int i = 0; i < requestData.size(); i++) {
                 int inputRecordIndex = batchNumber * batchSize + i;
                 com.skyflow.generated.rest.types.V1FlowTokenizeRequestObject reqObj = requestData.get(i);
-                List<String> groupNames = reqObj.getTokenGroupNames().isPresent()
-                        ? reqObj.getTokenGroupNames().get() : new ArrayList<>();
-                int groupCount = groupNames.size();
+                int groupCount = reqObj.getTokenGroupNames().isPresent()
+                        ? reqObj.getTokenGroupNames().get().size() : 0;
+                int entriesToConsume = groupCount > 0 ? groupCount : 1;
 
-                // Consume exactly groupCount entries from the flat response for this record
                 TokenizeSuccess successEntry = null;
-                for (int g = 0; g < groupCount && flatIndex < flatList.size(); g++, flatIndex++) {
-                    com.skyflow.generated.rest.types.V1FlowTokenizeResponseObject obj = flatList.get(flatIndex);
+                for (int j = 0; j < entriesToConsume && flatIndex < flatList.size(); j++) {
+                    com.skyflow.generated.rest.types.V1FlowTokenizeResponseObject obj = flatList.get(flatIndex++);
                     Map<String, Object> props = obj.getAdditionalProperties();
-
                     Object value = obj.getValue().isPresent() ? obj.getValue().get()
                             : (props != null ? props.get("value") : null);
                     String tokenGroupName = props != null ? (String) props.get("tokenGroupName") : null;

v3/src/main/java/com/skyflow/utils/validations/Validations.java

@@ -248,21 +248,17 @@ public static void validateTokenizeRequest(TokenizeRequest request) throws Skyfl
                 throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyValueInTokenizeRecord.getMessage());
             }
             List<String> tokenGroupNames = record.getTokenGroupNames();
-            if (tokenGroupNames == null || tokenGroupNames.isEmpty()) {
-                LogUtil.printErrorLog(Utils.parameterizedString(
-                        ErrorLogs.EMPTY_TOKEN_GROUP_NAMES_IN_TOKENIZE_RECORD.getLog(), InterfaceName.TOKENIZE.getName()
-                ));
-                throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyTokenGroupNamesInTokenizeRecord.getMessage());
-            }
-            for (int j = 0; j < tokenGroupNames.size(); j++) {
-                String groupName = tokenGroupNames.get(j);
-                if (groupName == null || groupName.trim().isEmpty()) {
-                    LogUtil.printErrorLog(Utils.parameterizedString(
-                            ErrorLogs.EMPTY_TOKEN_GROUP_NAME_IN_TOKENIZE_RECORD.getLog(),
-                            InterfaceName.TOKENIZE.getName(),
-                            String.valueOf(j)
-                    ));
-                    throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyTokenGroupNameInTokenizeRecord.getMessage());
+            if (tokenGroupNames != null) {
+                for (int j = 0; j < tokenGroupNames.size(); j++) {
+                    String groupName = tokenGroupNames.get(j);
+                    if (groupName == null || groupName.trim().isEmpty()) {
+                        LogUtil.printErrorLog(Utils.parameterizedString(
+                                ErrorLogs.EMPTY_TOKEN_GROUP_NAME_IN_TOKENIZE_RECORD.getLog(),
+                                InterfaceName.TOKENIZE.getName(),
+                                String.valueOf(j)
+                        ));
+                        throw new SkyflowException(ErrorCode.INVALID_INPUT.getCode(), ErrorMessage.EmptyTokenGroupNameInTokenizeRecord.getMessage());
+                    }
                 }
             }
         }

v3/src/test/java/com/skyflow/vault/controller/VaultControllerTokenizeTests.java

@@ -164,20 +164,6 @@ public void testValidation_emptyValue() {
         }
     }
 
-    @Test
-    public void testValidation_emptyTokenGroupNames() {
-        ArrayList<TokenizeRecord> records = new ArrayList<>();
-        records.add(TokenizeRecord.builder().value("val").tokenGroupNames(new ArrayList<>()).build());
-        TokenizeRequest req = TokenizeRequest.builder().data(records).build();
-        try {
-            Validations.validateTokenizeRequest(req);
-            fail("Expected SkyflowException for empty token group names");
-        } catch (SkyflowException e) {
-            assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
-            assertEquals(ErrorMessage.EmptyTokenGroupNamesInTokenizeRecord.getMessage(), e.getMessage());
-        }
-    }
-
     @Test
     public void testValidation_emptyTokenGroupNameInList() {
         ArrayList<TokenizeRecord> records = new ArrayList<>();

v3/src/test/java/com/skyflow/vault/data/TokenizeRequestTests.java

@@ -163,24 +163,21 @@ public void testValidateTokenizeRequest_nullTokenGroupNames() {
         TokenizeRequest req = TokenizeRequest.builder().data(records).build();
         try {
             Validations.validateTokenizeRequest(req);
-            Assert.fail("Expected SkyflowException for null tokenGroupNames");
         } catch (SkyflowException e) {
-            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
-            Assert.assertEquals(ErrorMessage.EmptyTokenGroupNamesInTokenizeRecord.getMessage(), e.getMessage());
+            Assert.fail("Should not throw for null tokenGroupNames: " + e.getMessage());
         }
     }
 
     @Test
     public void testValidateTokenizeRequest_emptyTokenGroupNames() {
+        // empty tokenGroupNames list is valid — the API tokenizes against all token groups
         ArrayList<TokenizeRecord> records = new ArrayList<>();
         records.add(TokenizeRecord.builder().value("val").tokenGroupNames(new ArrayList<>()).build());
         TokenizeRequest req = TokenizeRequest.builder().data(records).build();
         try {
             Validations.validateTokenizeRequest(req);
-            Assert.fail("Expected SkyflowException for empty tokenGroupNames");
         } catch (SkyflowException e) {
-            Assert.assertEquals(ErrorCode.INVALID_INPUT.getCode(), e.getHttpCode());
-            Assert.assertEquals(ErrorMessage.EmptyTokenGroupNamesInTokenizeRecord.getMessage(), e.getMessage());
+            Assert.fail("Should not throw for empty tokenGroupNames: " + e.getMessage());
         }
     }