diff --git a/codeflash/languages/java/remove_asserts.py b/codeflash/languages/java/remove_asserts.py index 0a84b687c..a5a5986c9 100644 --- a/codeflash/languages/java/remove_asserts.py +++ b/codeflash/languages/java/remove_asserts.py @@ -956,8 +956,29 @@ def _infer_type_from_assertion_args(self, original_text: str, method: str) -> st elif args_str.endswith(")"): args_str = args_str[:-1] - # Fast-path: only extract the first top-level argument instead of splitting all arguments. - first_arg = self._extract_first_arg(args_str) + # Fast-path: try to cheaply obtain the first top-level argument without invoking + # the full parser. If the first comma occurs before any special characters that + # would affect top-levelness (quotes/parens/braces), we can take the substring + # up to that comma as the first argument. + if not args_str: + return "Object" + + # Find first comma; if none, the entire args_str is the single argument + comma_idx = args_str.find(",") + if comma_idx == -1: + first_arg = args_str + else: + # If there are no special delimiter characters before this comma, we can + # safely take the substring as the first argument. + before = args_str[:comma_idx] + if not self._special_re.search(before): + first_arg = before + else: + # Fallback: use the full extractor which respects nesting/strings/generics. + first_arg = self._extract_first_arg(args_str) + if first_arg is None: + return "Object" + if not first_arg: return "Object"