fix: auto-add missing standard library imports in AI-generated Java tests

AI-generated test code sometimes uses standard library classes like
Arrays, List, HashMap etc. without the corresponding import statement,
causing compilation failures. Add ensure_common_java_imports() that
detects usage of common classes and adds missing imports automatically.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: mohamedashrraf222

codeflash/languages/java/instrumentation.py

@@ -1260,6 +1260,35 @@ def remove_instrumentation(source: str) -> str:
     return source
 
 
+_COMMON_JAVA_IMPORTS = {
+    "Arrays": "import java.util.Arrays;",
+    "List": "import java.util.List;",
+    "ArrayList": "import java.util.ArrayList;",
+    "Map": "import java.util.Map;",
+    "HashMap": "import java.util.HashMap;",
+    "Set": "import java.util.Set;",
+    "HashSet": "import java.util.HashSet;",
+    "Collections": "import java.util.Collections;",
+    "Collectors": "import java.util.stream.Collectors;",
+    "Random": "import java.util.Random;",
+    "BigDecimal": "import java.math.BigDecimal;",
+    "BigInteger": "import java.math.BigInteger;",
+}
+
+
+def ensure_common_java_imports(test_code: str) -> str:
+    for class_name, import_stmt in _COMMON_JAVA_IMPORTS.items():
+        if not re.search(rf"\b{class_name}\b", test_code):
+            continue
+        if import_stmt in test_code:
+            continue
+        package = import_stmt.split()[1].rsplit(".", 1)[0]
+        if f"import {package}.*;" in test_code:
+            continue
+        test_code = _add_import(test_code, import_stmt)
+    return test_code
+
+
 def instrument_generated_java_test(
     test_code: str,
     function_name: str,
@@ -1290,6 +1319,7 @@ def instrument_generated_java_test(
     from codeflash.languages.java.remove_asserts import transform_java_assertions
 
     test_code = transform_java_assertions(test_code, function_name, qualified_name)
+    test_code = ensure_common_java_imports(test_code)
 
     # Extract class name from the test code
     # Use pattern that starts at beginning of line to avoid matching words in comments