Skip to content

Commit ccf3441

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
chore: chore
PiperOrigin-RevId: 897278822
1 parent 8e8146a commit ccf3441

3 files changed

Lines changed: 124 additions & 11 deletions

File tree

‎src/main/java/com/google/genai/ReplayApiClient.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,16 @@ private static Object normalizeKeyCase(Object obj) {
352352
entry -> {
353353
String newKey = Common.snakeToCamel(entry.getKey());
354354
Object normalizedValue = normalizeKeyCase(entry.getValue());
355+
356+
if (newKey.equals("data") || newKey.equals("imageBytes")) {
357+
if (normalizedValue instanceof JsonNode && ((JsonNode) normalizedValue).isTextual()) {
358+
String base64Str = ((JsonNode) normalizedValue).asText();
359+
// Normalize URL-safe Base64 to Standard Base64 and remove padding for comparison
360+
base64Str = base64Str.replace('-', '+').replace('_', '/').replaceAll("=", "");
361+
normalizedValue = JsonSerializable.toJsonNode(base64Str);
362+
}
363+
}
364+
355365
normalizedNode.set(newKey, JsonSerializable.toJsonNode(normalizedValue));
356366
});
357367
return normalizedNode;

‎src/test/java/com/google/genai/TableTest.java‎

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -194,16 +194,6 @@ private static Collection<DynamicTest> createTestCases(
194194
String msg = " => Test skipped: parameters contain unsupported union type";
195195
return Collections.singletonList(DynamicTest.dynamicTest(testName + msg, () -> {}));
196196
}
197-
// Edit image ReferenceImages are not correctly deserialized for replay tests
198-
if (testName.contains("models.edit_image")
199-
|| testName.contains("batches.create.test_with_image_blob")) { // TODO(b/431798111)
200-
String msg = " => Test skipped: replay tests are not supported for edit_image";
201-
return Collections.singletonList(DynamicTest.dynamicTest(testName + msg, () -> {}));
202-
}
203-
if (testName.contains("models.embed_content.test_new_api_inline_pdf")) {
204-
String msg = " => Test skipped: inline byte deserialization fails";
205-
return Collections.singletonList(DynamicTest.dynamicTest(testName + msg, () -> {}));
206-
}
207197
// TODO(b/457846189): Support models.list filter parameter
208198
if (testName.contains("models.list.test_tuned_models_with_filter")
209199
|| testName.contains("models.list.test_tuned_models.vertex")) {
@@ -224,7 +214,8 @@ private static Collection<DynamicTest> createTestCases(
224214
Object fromValue = fromParameters.getOrDefault(parameterName, null);
225215
// May throw IllegalArgumentException here
226216
Object parameter =
227-
JsonSerializable.objectMapper.convertValue(fromValue, method.getParameterTypes()[i]);
217+
TestUtils.getTestObjectMapper().convertValue(
218+
fromValue, TestUtils.getTestObjectMapper().constructType(method.getGenericParameterTypes()[i]));
228219
if (method.getName().equals("embedContent") && parameter instanceof List) {
229220
throw new IllegalArgumentException();
230221
}
@@ -337,6 +328,38 @@ private static Map<String, Object> prepareParameters(TestTableItem testTableItem
337328
"source.scribbleImage.image.imageBytes",
338329
new ReplayBase64Sanitizer(),
339330
false);
331+
ReplaySanitizer.sanitizeMapByPath(
332+
fromParameters,
333+
"[]referenceImages.referenceImage.imageBytes",
334+
new ReplayBase64Sanitizer(),
335+
false);
336+
ReplaySanitizer.sanitizeMapByPath(
337+
fromParameters,
338+
"referenceImages.[]referenceImage.imageBytes",
339+
new ReplayBase64Sanitizer(),
340+
false);
341+
ReplaySanitizer.sanitizeMapByPath(
342+
fromParameters, "[]contents.[]parts.inlineData.data", new ReplayBase64Sanitizer(), false);
343+
ReplaySanitizer.sanitizeMapByPath(
344+
fromParameters,
345+
"contents.[]parts.inlineData.data",
346+
new ReplayBase64Sanitizer(),
347+
false);
348+
ReplaySanitizer.sanitizeMapByPath(
349+
fromParameters,
350+
"[]contents.[]parts.functionResponse.[]parts.inlineData.data",
351+
new ReplayBase64Sanitizer(),
352+
false);
353+
ReplaySanitizer.sanitizeMapByPath(
354+
fromParameters,
355+
"contents.[]parts.functionResponse.[]parts.inlineData.data",
356+
new ReplayBase64Sanitizer(),
357+
false);
358+
ReplaySanitizer.sanitizeMapByPath(
359+
fromParameters,
360+
"src.[]inlinedRequests.[]contents.[]parts.inlineData.data",
361+
new ReplayBase64Sanitizer(),
362+
false);
340363
return fromParameters;
341364
}
342365

‎src/test/java/com/google/genai/TestUtils.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,91 @@
1616

1717
package com.google.genai;
1818

19+
import com.fasterxml.jackson.core.JsonParser;
20+
import com.fasterxml.jackson.databind.DeserializationContext;
21+
import com.fasterxml.jackson.databind.JsonNode;
22+
import com.fasterxml.jackson.databind.ObjectMapper;
23+
import com.fasterxml.jackson.databind.deser.std.StdDeserializer;
24+
import com.fasterxml.jackson.databind.module.SimpleModule;
25+
import com.google.genai.types.ContentReferenceImage;
26+
import com.google.genai.types.ControlReferenceImage;
27+
import com.google.genai.types.MaskReferenceImage;
28+
import com.google.genai.types.RawReferenceImage;
29+
import com.google.genai.types.ReferenceImage;
30+
import com.google.genai.types.StyleReferenceImage;
31+
import com.google.genai.types.SubjectReferenceImage;
32+
import java.io.IOException;
33+
1934
public final class TestUtils {
2035
static final String API_KEY = "api-key";
2136
static final String PROJECT = "project";
2237
static final String LOCATION = "location";
2338

39+
private static ObjectMapper testObjectMapper;
40+
41+
public static ObjectMapper getTestObjectMapper() {
42+
if (testObjectMapper == null) {
43+
testObjectMapper = JsonSerializable.objectMapper.copy();
44+
SimpleModule customModule = new SimpleModule();
45+
customModule.addDeserializer(ReferenceImage.class, new ReferenceImageDeserializer());
46+
testObjectMapper.registerModule(customModule);
47+
}
48+
return testObjectMapper;
49+
}
50+
51+
private static class ReferenceImageDeserializer extends StdDeserializer<ReferenceImage> {
52+
public ReferenceImageDeserializer() {
53+
this(null);
54+
}
55+
56+
public ReferenceImageDeserializer(Class<?> vc) {
57+
super(vc);
58+
}
59+
60+
@Override
61+
public ReferenceImage deserialize(JsonParser jp, DeserializationContext ctxt)
62+
throws IOException {
63+
JsonNode node = jp.getCodec().readTree(jp);
64+
if (node.isObject()) {
65+
com.fasterxml.jackson.databind.node.ObjectNode objNode =
66+
(com.fasterxml.jackson.databind.node.ObjectNode) node;
67+
if (objNode.has("maskImageConfig")) {
68+
objNode.set("config", objNode.get("maskImageConfig"));
69+
}
70+
if (objNode.has("styleImageConfig")) {
71+
objNode.set("config", objNode.get("styleImageConfig"));
72+
}
73+
if (objNode.has("controlImageConfig")) {
74+
objNode.set("config", objNode.get("controlImageConfig"));
75+
}
76+
if (objNode.has("subjectImageConfig")) {
77+
objNode.set("config", objNode.get("subjectImageConfig"));
78+
}
79+
if (objNode.has("contentImageConfig")) {
80+
objNode.set("config", objNode.get("contentImageConfig"));
81+
}
82+
}
83+
84+
if (node.has("referenceType")) {
85+
String type = node.get("referenceType").asText();
86+
if ("REFERENCE_TYPE_RAW".equals(type)) {
87+
return jp.getCodec().treeToValue(node, RawReferenceImage.class);
88+
} else if ("REFERENCE_TYPE_MASK".equals(type)) {
89+
return jp.getCodec().treeToValue(node, MaskReferenceImage.class);
90+
} else if ("REFERENCE_TYPE_CONTROL".equals(type)) {
91+
return jp.getCodec().treeToValue(node, ControlReferenceImage.class);
92+
} else if ("REFERENCE_TYPE_STYLE".equals(type)) {
93+
return jp.getCodec().treeToValue(node, StyleReferenceImage.class);
94+
} else if ("REFERENCE_TYPE_SUBJECT".equals(type)) {
95+
return jp.getCodec().treeToValue(node, SubjectReferenceImage.class);
96+
} else if ("REFERENCE_TYPE_CONTENT".equals(type)) {
97+
return jp.getCodec().treeToValue(node, ContentReferenceImage.class);
98+
}
99+
}
100+
throw new IOException("Unknown or missing referenceType for ReferenceImage");
101+
}
102+
}
103+
24104
private TestUtils() {}
25105

26106
/** Creates a client given the vertexAI and replayId. Can be used in replay tests. */

0 commit comments

Comments
 (0)