Async tests.

newrelic · Dec 13, 2023 · 8c97a19 · 8c97a19
1 parent 0b7b33e
commit 8c97a19
Showing 1 changed file with 117 additions and 16 deletions.
diff --git a/tests/mlmodel_openai/test_embeddings_error_v1.py b/tests/mlmodel_openai/test_embeddings_error_v1.py
@@ -64,11 +64,11 @@
         "user": {},
     },
 )
-# @validate_span_events(
-#     exact_agents={
-#         "error.message": "Unknown Prompt:\nModel does not exist.",
-#     }
-#)
+@validate_span_events(
+    exact_agents={
+        "error.message": "Unknown Prompt:\nModel does not exist.",
+    }
+)
 @validate_transaction_metrics(
     name="test_embeddings_error_v1:test_embeddings_invalid_request_error_invalid_model",
     scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
@@ -119,11 +119,11 @@ def test_embeddings_invalid_request_error_invalid_model(set_trace_info, sync_ope
         "user": {},
     },
 )
-# @validate_span_events(
-#     exact_agents={
-#         "error.message": "Incorrect API key provided: DEADBEEF. You can find your API key at https://platform.openai.com/account/api-keys.",
-#     }
-# )
+@validate_span_events(
+    exact_agents={
+        "error.message": "Incorrect API key provided: DEADBEEF. You can find your API key at https://platform.openai.com/account/api-keys.",
+    }
+)
 @validate_transaction_metrics(
     name="test_embeddings_error_v1:test_embeddings_wrong_api_key_error",
     scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
@@ -136,7 +136,7 @@ def test_embeddings_invalid_request_error_invalid_model(set_trace_info, sync_ope
 @validate_custom_events(embedding_invalid_key_error_events)
 @validate_custom_event_count(count=1)
 @background_task()
-def test_embeddings_wrong_api_key_error(monkeypatch, set_trace_info):
+def test_embeddings_wrong_api_key_error(set_trace_info):
     with pytest.raises(openai.AuthenticationError):
         set_trace_info()
         wrong_api_key_client = openai.OpenAI(api_key="DEADBEEF")
@@ -175,11 +175,11 @@ def test_embeddings_wrong_api_key_error(monkeypatch, set_trace_info):
         "user": {},
     },
 )
-# @validate_span_events(
-#     exact_agents={
-#         "error.message": "Connection error.",
-#     }
-# )
+@validate_span_events(
+    exact_agents={
+        "error.message": "Connection error.",
+    }
+)
 @validate_transaction_metrics(
     name="test_embeddings_error_v1:test_embeddings_authentication_error",
     scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
@@ -199,4 +199,105 @@ def test_embeddings_authentication_error(set_trace_info):
         no_api_key_client.embeddings.create(input="No API key.", model="text-embedding-ada-002")
 
 
+@dt_enabled
+@reset_core_stats_engine()
+@validate_error_trace_attributes(
+    callable_name(openai.InternalServerError),
+    exact_attrs={
+        "agent": {},
+        "intrinsic": {},
+        "user": {},
+    },
+)
+@validate_span_events(
+    exact_agents={
+        "error.message": "Unknown Prompt:\nModel does not exist.",
+    }
+)
+@validate_transaction_metrics(
+    name="test_embeddings_error_v1:test_embeddings_invalid_request_error_invalid_model_async",
+    scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    rollup_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
+    background_task=True,
+)
+@validate_custom_events(invalid_model_events)
+@validate_custom_event_count(count=1)
+@background_task()
+def test_embeddings_invalid_request_error_invalid_model_async(set_trace_info, async_openai_client, loop):
+    with pytest.raises(openai.InternalServerError):
+        set_trace_info()
+        loop.run_until_complete(async_openai_client.embeddings.create(input="Model does not exist.", model="does-not-exist"))
+
+
+@dt_enabled
+@reset_core_stats_engine()
+@validate_error_trace_attributes(
+    callable_name(openai.AuthenticationError),
+    exact_attrs={
+        "agent": {},
+        "intrinsic": {},
+        "user": {},
+    },
+)
+@validate_span_events(
+    exact_agents={
+        "error.message": "Incorrect API key provided: DEADBEEF. You can find your API key at https://platform.openai.com/account/api-keys.",
+    }
+)
+@validate_transaction_metrics(
+    name="test_embeddings_error_v1:test_embeddings_wrong_api_key_error_async",
+    scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    rollup_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
+    background_task=True,
+)
+@validate_custom_events(embedding_invalid_key_error_events)
+@validate_custom_event_count(count=1)
+@background_task()
+def test_embeddings_wrong_api_key_error_async(set_trace_info, loop):
+    with pytest.raises(openai.AuthenticationError):
+        set_trace_info()
+        wrong_api_key_client = openai.AsyncOpenAI(api_key="DEADBEEF")
+        loop.run_until_complete(wrong_api_key_client.embeddings.create(input="Embedded: Invalid API key.", model="text-embedding-ada-002"))
+
+
+
+@dt_enabled
+@reset_core_stats_engine()
+@validate_error_trace_attributes(
+    callable_name(openai.APIConnectionError),
+    exact_attrs={
+        "agent": {},
+        "intrinsic": {},
+        "user": {},
+    },
+)
+@validate_span_events(
+    exact_agents={
+        "error.message": "Connection error.",
+    }
+)
+@validate_transaction_metrics(
+    name="test_embeddings_error_v1:test_embeddings_authentication_error_async",
+    scoped_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    rollup_metrics=[("Llm/embedding/OpenAI/create", 1)],
+    custom_metrics=[
+        ("Python/ML/OpenAI/%s" % openai.__version__, 1),
+    ],
+    background_task=True,
+)
+@validate_custom_events(embedding_auth_error_events)
+@validate_custom_event_count(count=1)
+@background_task()
+def test_embeddings_authentication_error_async(set_trace_info, loop):
+    with pytest.raises(openai.APIConnectionError):
+        set_trace_info()
+        no_api_key_client = openai.AsyncOpenAI(api_key="")
+        loop.run_until_complete(no_api_key_client.embeddings.create(input="No API key.", model="text-embedding-ada-002"))
+