Skip to content

Commit d6d2433

Browse files
authored
Fix python postcommit 30513 (#37849)
* Fix Vertex AI image embedding test * Removed unused import
1 parent b92c436 commit d6d2433

1 file changed

Lines changed: 2 additions & 7 deletions

File tree

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

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616

1717
"""Tests for apache_beam.ml.rag.embeddings.vertex_ai."""
1818

19-
import os
2019
import shutil
2120
import struct
2221
import tempfile
@@ -186,21 +185,17 @@ def test_image_embedding_pipeline(self):
186185

187186
def test_image_embedding_pipeline_from_path(self):
188187
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-
193188
test_items = [
194189
EmbeddableItem.from_image(
195-
img_path, id="img1", metadata={"source": "test"}),
190+
png_bytes, id="img1", metadata={"source": "test"}),
196191
]
197192

198193
expected = [
199194
EmbeddableItem(
200195
id="img1",
201196
embedding=Embedding(dense_embedding=[0.0] * 1408),
202197
metadata={"source": "test"},
203-
content=Content(image=img_path)),
198+
content=Content(image=png_bytes)),
204199
]
205200

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

0 commit comments

Comments
 (0)