Skip to content

Commit ba77742

Browse files
authored
[Backport 2.9] Add null key check for remote inference integration tests (#1136)
* Add null key check for remote inference integration tests Signed-off-by: Ryan Bogan <[email protected]> * Fix spotless Signed-off-by: Ryan Bogan <[email protected]> --------- Signed-off-by: Ryan Bogan <[email protected]>
1 parent 51c0ccd commit ba77742

File tree

1 file changed

+38
-7
lines changed

1 file changed

+38
-7
lines changed

plugin/src/test/java/org/opensearch/ml/rest/RestMLRemoteInferenceIT.java

Lines changed: 38 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424

2525
public class RestMLRemoteInferenceIT extends MLCommonsRestTestCase {
2626

27+
private final String OPENAI_KEY = System.getenv("OPENAI_KEY");
28+
private final String COHERE_KEY = System.getenv("COHERE_KEY");
29+
2730
private final String completionModelConnectorEntity = "{\n"
2831
+ "\"name\": \"OpenAI Connector\",\n"
2932
+ "\"description\": \"The connector to public OpenAI model service for GPT 3.5\",\n"
@@ -39,7 +42,7 @@ public class RestMLRemoteInferenceIT extends MLCommonsRestTestCase {
3942
+ " },\n"
4043
+ " \"credential\": {\n"
4144
+ " \"openAI_key\": \""
42-
+ System.getenv("OPENAI_KEY")
45+
+ OPENAI_KEY
4346
+ "\"\n"
4447
+ " },\n"
4548
+ " \"actions\": [\n"
@@ -133,6 +136,10 @@ public void testDeployRemoteModel() throws IOException, InterruptedException {
133136

134137
@Ignore
135138
public void testPredictRemoteModel() throws IOException, InterruptedException {
139+
// Skip test if key is null
140+
if (OPENAI_KEY == null) {
141+
return;
142+
}
136143
Response response = createConnector(completionModelConnectorEntity);
137144
Map responseMap = parseResponseToMap(response);
138145
String connectorId = (String) responseMap.get("connector_id");
@@ -186,6 +193,10 @@ public void testUndeployRemoteModel() throws IOException, InterruptedException {
186193

187194
@Ignore
188195
public void testOpenAIChatCompletionModel() throws IOException, InterruptedException {
196+
// Skip test if key is null
197+
if (OPENAI_KEY == null) {
198+
return;
199+
}
189200
String entity = "{\n"
190201
+ " \"name\": \"OpenAI chat model Connector\",\n"
191202
+ " \"description\": \"The connector to public OpenAI model service for GPT 3.5\",\n"
@@ -201,7 +212,7 @@ public void testOpenAIChatCompletionModel() throws IOException, InterruptedExcep
201212
+ " },\n"
202213
+ " \"credential\": {\n"
203214
+ " \"openAI_key\": \""
204-
+ System.getenv("OPENAI_KEY")
215+
+ OPENAI_KEY
205216
+ "\"\n"
206217
+ " },\n"
207218
+ " \"actions\": [\n"
@@ -243,6 +254,10 @@ public void testOpenAIChatCompletionModel() throws IOException, InterruptedExcep
243254

244255
@Ignore
245256
public void testOpenAIEditsModel() throws IOException, InterruptedException {
257+
// Skip test if key is null
258+
if (OPENAI_KEY == null) {
259+
return;
260+
}
246261
String entity = "{\n"
247262
+ " \"name\": \"OpenAI Edit model Connector\",\n"
248263
+ " \"description\": \"The connector to public OpenAI edit model service\",\n"
@@ -256,7 +271,7 @@ public void testOpenAIEditsModel() throws IOException, InterruptedException {
256271
+ " },\n"
257272
+ " \"credential\": {\n"
258273
+ " \"openAI_key\": \""
259-
+ System.getenv("OPENAI_KEY")
274+
+ OPENAI_KEY
260275
+ "\"\n"
261276
+ " },\n"
262277
+ " \"actions\": [\n"
@@ -309,6 +324,10 @@ public void testOpenAIEditsModel() throws IOException, InterruptedException {
309324

310325
@Ignore
311326
public void testOpenAIModerationsModel() throws IOException, InterruptedException {
327+
// Skip test if key is null
328+
if (OPENAI_KEY == null) {
329+
return;
330+
}
312331
String entity = "{\n"
313332
+ " \"name\": \"OpenAI moderations model Connector\",\n"
314333
+ " \"description\": \"The connector to public OpenAI moderations model service\",\n"
@@ -322,7 +341,7 @@ public void testOpenAIModerationsModel() throws IOException, InterruptedExceptio
322341
+ " },\n"
323342
+ " \"credential\": {\n"
324343
+ " \"openAI_key\": \""
325-
+ System.getenv("OPENAI_KEY")
344+
+ OPENAI_KEY
326345
+ "\"\n"
327346
+ " },\n"
328347
+ " \"actions\": [\n"
@@ -372,6 +391,10 @@ public void testOpenAIModerationsModel() throws IOException, InterruptedExceptio
372391

373392
@Ignore
374393
public void testOpenAITextEmbeddingModel() throws IOException, InterruptedException {
394+
// Skip test if key is null
395+
if (OPENAI_KEY == null) {
396+
return;
397+
}
375398
String entity = "{\n"
376399
+ " \"name\": \"OpenAI text embedding model Connector\",\n"
377400
+ " \"description\": \"The connector to public OpenAI text embedding model service\",\n"
@@ -385,7 +408,7 @@ public void testOpenAITextEmbeddingModel() throws IOException, InterruptedExcept
385408
+ " },\n"
386409
+ " \"credential\": {\n"
387410
+ " \"openAI_key\": \""
388-
+ System.getenv("OPENAI_KEY")
411+
+ OPENAI_KEY
389412
+ "\"\n"
390413
+ " },\n"
391414
+ " \"actions\": [\n"
@@ -430,6 +453,10 @@ public void testOpenAITextEmbeddingModel() throws IOException, InterruptedExcept
430453
}
431454

432455
public void testCohereGenerateTextModel() throws IOException, InterruptedException {
456+
// Skip test if key is null
457+
if (COHERE_KEY == null) {
458+
return;
459+
}
433460
String entity = "{\n"
434461
+ " \"name\": \"Cohere generate text model Connector\",\n"
435462
+ " \"description\": \"The connector to public Cohere generate text model service\",\n"
@@ -443,7 +470,7 @@ public void testCohereGenerateTextModel() throws IOException, InterruptedExcepti
443470
+ " },\n"
444471
+ " \"credential\": {\n"
445472
+ " \"cohere_key\": \""
446-
+ System.getenv("COHERE_KEY")
473+
+ COHERE_KEY
447474
+ "\"\n"
448475
+ " },\n"
449476
+ " \"actions\": [\n"
@@ -491,6 +518,10 @@ public void testCohereGenerateTextModel() throws IOException, InterruptedExcepti
491518
}
492519

493520
public void testCohereClassifyModel() throws IOException, InterruptedException {
521+
// Skip test if key is null
522+
if (COHERE_KEY == null) {
523+
return;
524+
}
494525
String entity = "{\n"
495526
+ " \"name\": \"Cohere classify model Connector\",\n"
496527
+ " \"description\": \"The connector to public Cohere classify model service\",\n"
@@ -504,7 +535,7 @@ public void testCohereClassifyModel() throws IOException, InterruptedException {
504535
+ " },\n"
505536
+ " \"credential\": {\n"
506537
+ " \"cohere_key\": \""
507-
+ System.getenv("COHERE_KEY")
538+
+ COHERE_KEY
508539
+ "\"\n"
509540
+ " },\n"
510541
+ " \"actions\": [\n"

0 commit comments

Comments
 (0)