Skip to content

Commit 6e28454

Browse files
committed
Fix Vertex AI image embedding test
1 parent d6bc507 commit 6e28454

1 file changed

Lines changed: 2 additions & 6 deletions

File tree

sdks/python/apache_beam/ml/rag/embeddings/vertex_ai_test.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -186,21 +186,17 @@ def test_image_embedding_pipeline(self):
186186

187187
def test_image_embedding_pipeline_from_path(self):
188188
png_bytes = _create_png_bytes()
189-
img_path = os.path.join(self.artifact_location, 'test_img.png')
190-
with open(img_path, 'wb') as f:
191-
f.write(png_bytes)
192-
193189
test_items = [
194190
EmbeddableItem.from_image(
195-
img_path, id="img1", metadata={"source": "test"}),
191+
png_bytes, id="img1", metadata={"source": "test"}),
196192
]
197193

198194
expected = [
199195
EmbeddableItem(
200196
id="img1",
201197
embedding=Embedding(dense_embedding=[0.0] * 1408),
202198
metadata={"source": "test"},
203-
content=Content(image=img_path)),
199+
content=Content(image=png_bytes)),
204200
]
205201

206202
artifact_location = tempfile.mkdtemp(prefix='vertex_ai_img_path_')

0 commit comments

Comments
 (0)