Skip to content

Commit 0f410ac

Browse files
sjrlsjrl
authored
ci: Reduce impact of integration tests for Watsonx integration (#2099)
* Reduce integration tests * Further reduce token usage * Formatting
1 parent 63bbbaa commit 0f410ac

4 files changed

Lines changed: 18 additions & 110 deletions

File tree

‎integrations/watsonx/tests/test_chat_generator.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -393,7 +393,7 @@ def test_live_run(self):
393393
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
394394
generation_kwargs={"max_tokens": 50, "temperature": 0.7, "top_p": 0.9},
395395
)
396-
messages = [ChatMessage.from_user("What's the capital of France?")]
396+
messages = [ChatMessage.from_user("What's the capital of France? Answer concisely.")]
397397
results = generator.run(messages=messages)
398398

399399
assert isinstance(results, dict)
@@ -417,7 +417,7 @@ def test_live_run_streaming(self):
417417
def callback(chunk: StreamingChunk):
418418
collected_chunks.append(chunk)
419419

420-
messages = [ChatMessage.from_user("Explain quantum computing")]
420+
messages = [ChatMessage.from_user("What's the capital of France? Answer concisely.")]
421421
results = generator.run(messages=messages, streaming_callback=callback)
422422

423423
assert isinstance(results, dict)
@@ -437,7 +437,7 @@ async def test_live_run_async(self):
437437
generator = WatsonxChatGenerator(
438438
model="ibm/granite-3-2b-instruct", project_id=Secret.from_env_var("WATSONX_PROJECT_ID")
439439
)
440-
messages = [ChatMessage.from_user("What's the capital of Germany?")]
440+
messages = [ChatMessage.from_user("What's the capital of Germany? Answer concisely.")]
441441
results = await generator.run_async(messages=messages)
442442

443443
assert isinstance(results, dict)

‎integrations/watsonx/tests/test_document_embedder.py‎

Lines changed: 6 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -217,18 +217,16 @@ def test_run(self, test_documents):
217217
)
218218
def test_batch_processing(self, test_documents):
219219
"""Test that batch processing works"""
220-
many_documents = test_documents * 50
221-
222220
embedder = WatsonxDocumentEmbedder(
223221
model="ibm/slate-30m-english-rtrvr",
224222
api_key=Secret.from_env_var("WATSONX_API_KEY"),
225223
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
226-
batch_size=50,
224+
batch_size=2,
227225
truncate_input_tokens=128,
228226
)
229227

230-
result = embedder.run(many_documents)
231-
assert len(result["documents"]) == 150
228+
result = embedder.run(test_documents)
229+
assert len(result["documents"]) == 3
232230
assert all(doc.embedding is not None for doc in result["documents"])
233231

234232
@pytest.mark.skipif(
@@ -237,51 +235,16 @@ def test_batch_processing(self, test_documents):
237235
)
238236
def test_text_truncation(self):
239237
"""Test that truncation works with long documents"""
240-
long_content = "This is a very long document. " * 1000
238+
long_content = "This is a very long document. " * 10
241239
long_document = Document(content=long_content)
242240

243241
embedder = WatsonxDocumentEmbedder(
244242
model="ibm/slate-30m-english-rtrvr",
245243
api_key=Secret.from_env_var("WATSONX_API_KEY"),
246244
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
247-
truncate_input_tokens=128,
245+
truncate_input_tokens=4,
248246
)
249247

250248
result = embedder.run([long_document])
251249
assert len(result["documents"][0].embedding) > 0
252-
assert result["meta"]["truncate_input_tokens"] == 128
253-
254-
@pytest.mark.skipif(
255-
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
256-
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
257-
)
258-
def test_prefix_suffix(self, test_documents):
259-
"""Test that prefix and suffix are correctly applied"""
260-
embedder = WatsonxDocumentEmbedder(
261-
model="ibm/slate-30m-english-rtrvr",
262-
api_key=Secret.from_env_var("WATSONX_API_KEY"),
263-
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
264-
prefix="PREFIX: ",
265-
suffix=" :SUFFIX",
266-
truncate_input_tokens=128,
267-
)
268-
269-
result = embedder.run([test_documents[0]])
270-
assert result["documents"][0].embedding is not None
271-
272-
@pytest.mark.skipif(
273-
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
274-
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
275-
)
276-
def test_concurrency_handling(self, test_documents):
277-
"""Test that concurrency limits are respected"""
278-
embedder = WatsonxDocumentEmbedder(
279-
model="ibm/slate-30m-english-rtrvr",
280-
api_key=Secret.from_env_var("WATSONX_API_KEY"),
281-
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
282-
concurrency_limit=2,
283-
batch_size=1,
284-
)
285-
286-
result = embedder.run(test_documents)
287-
assert len(result["documents"]) == 3
250+
assert result["meta"]["truncate_input_tokens"] == 4

‎integrations/watsonx/tests/test_generator.py‎

Lines changed: 6 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -370,7 +370,10 @@ def test_live_run(self):
370370
generation_kwargs={"max_tokens": 50, "temperature": 0.7, "top_p": 0.9},
371371
)
372372

373-
result = generator.run(prompt="What's the capital of France?")
373+
result = generator.run(
374+
prompt="What's the capital of France? Answer concisely.",
375+
system_prompt="You are a geography expert. Answer concisely.",
376+
)
374377

375378
assert isinstance(result, dict)
376379
assert "replies" in result
@@ -383,27 +386,6 @@ def test_live_run(self):
383386
assert len(result["replies"][0]) > 0
384387
assert isinstance(result["meta"][0], dict)
385388

386-
@pytest.mark.skipif(
387-
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
388-
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
389-
)
390-
def test_live_run_with_system_prompt(self):
391-
generator = WatsonxGenerator(
392-
model="ibm/granite-3-2b-instruct",
393-
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
394-
)
395-
396-
result = generator.run(
397-
prompt="What's the capital of France?", system_prompt="You are a geography expert. Answer concisely."
398-
)
399-
400-
assert isinstance(result, dict)
401-
assert "replies" in result
402-
assert "meta" in result
403-
assert len(result["replies"]) == 1
404-
assert isinstance(result["replies"][0], str)
405-
assert len(result["replies"][0]) > 0
406-
407389
@pytest.mark.skipif(
408390
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
409391
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
@@ -418,7 +400,7 @@ def test_live_run_streaming(self):
418400
def callback(chunk: StreamingChunk):
419401
collected_chunks.append(chunk)
420402

421-
result = generator.run(prompt="Explain quantum computing", streaming_callback=callback)
403+
result = generator.run(prompt="What is the capital of France? Answer concisely.", streaming_callback=callback)
422404

423405
assert isinstance(result, dict)
424406
assert "replies" in result
@@ -441,7 +423,7 @@ async def test_live_run_async(self):
441423
model="ibm/granite-3-2b-instruct", project_id=Secret.from_env_var("WATSONX_PROJECT_ID")
442424
)
443425

444-
result = await generator.run_async(prompt="What's the capital of Germany?")
426+
result = await generator.run_async(prompt="What's the capital of Germany? Answer concisely.")
445427

446428
assert isinstance(result, dict)
447429
assert "replies" in result

‎integrations/watsonx/tests/test_text_embedder.py‎

Lines changed: 3 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -174,43 +174,6 @@ def test_run(self):
174174
assert "slate" in result["meta"]["model"].lower()
175175
assert "30m" in result["meta"]["model"].lower()
176176

177-
@pytest.mark.integration
178-
@pytest.mark.skipif(
179-
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
180-
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
181-
)
182-
def test_long_text(self):
183-
"""Test with longer text input"""
184-
embedder = WatsonxTextEmbedder(
185-
model="ibm/slate-30m-english-rtrvr",
186-
api_key=Secret.from_env_var("WATSONX_API_KEY"),
187-
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
188-
truncate_input_tokens=128,
189-
)
190-
long_text = "This is a test text that should work fine. " * 10
191-
result = embedder.run(long_text)
192-
193-
assert len(result["embedding"]) > 0
194-
assert "model" in result["meta"]
195-
196-
@pytest.mark.integration
197-
@pytest.mark.skipif(
198-
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
199-
reason="WATSONX_API_KEY or WATSONX_PROJECT_ID not set",
200-
)
201-
def test_different_model(self):
202-
"""Test with a different model if available"""
203-
embedder = WatsonxTextEmbedder(
204-
model="ibm/slate-125m-english-rtrvr",
205-
api_key=Secret.from_env_var("WATSONX_API_KEY"),
206-
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
207-
truncate_input_tokens=128,
208-
)
209-
result = embedder.run("Test different model")
210-
211-
assert len(result["embedding"]) > 0
212-
assert "125m" in result["meta"]["model"]
213-
214177
@pytest.mark.integration
215178
@pytest.mark.skipif(
216179
not os.environ.get("WATSONX_API_KEY") or not os.environ.get("WATSONX_PROJECT_ID"),
@@ -244,11 +207,11 @@ def test_text_truncation(self):
244207
model="ibm/slate-30m-english-rtrvr",
245208
api_key=Secret.from_env_var("WATSONX_API_KEY"),
246209
project_id=Secret.from_env_var("WATSONX_PROJECT_ID"),
247-
truncate_input_tokens=128,
210+
truncate_input_tokens=4,
248211
)
249212

250-
long_text = "This is a test sentence. " * 1000
213+
long_text = "This is a test sentence. " * 10
251214
result = embedder.run(long_text)
252215

253216
assert len(result["embedding"]) > 0
254-
assert result["meta"]["truncated_input_tokens"] == 128
217+
assert result["meta"]["truncated_input_tokens"] == 4

0 commit comments

Comments
 (0)