From 65b231d40bfd687f8ec305fa806246f72b82e389 Mon Sep 17 00:00:00 2001
From: Erick Friis <erick@langchain.dev>
Date: Thu, 18 Jan 2024 09:45:44 -0800
Subject: [PATCH] mistralai[patch]: async integration tests (#16214)

---
 .../integration_tests/test_embeddings.py      | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/libs/partners/mistralai/tests/integration_tests/test_embeddings.py b/libs/partners/mistralai/tests/integration_tests/test_embeddings.py
index b179f64d80f..166bf5c5c15 100644
--- a/libs/partners/mistralai/tests/integration_tests/test_embeddings.py
+++ b/libs/partners/mistralai/tests/integration_tests/test_embeddings.py
@@ -17,3 +17,37 @@ def test_mistralai_embedding_query() -> None:
     embedding = MistralAIEmbeddings()
     output = embedding.embed_query(document)
     assert len(output) == 1024
+
+
+async def test_mistralai_embedding_documents_async() -> None:
+    """Test MistralAI embeddings for documents."""
+    documents = ["foo bar", "test document"]
+    embedding = MistralAIEmbeddings()
+    output = await embedding.aembed_documents(documents)
+    assert len(output) == 2
+    assert len(output[0]) == 1024
+
+
+async def test_mistralai_embedding_query_async() -> None:
+    """Test MistralAI embeddings for query."""
+    document = "foo bar"
+    embedding = MistralAIEmbeddings()
+    output = await embedding.aembed_query(document)
+    assert len(output) == 1024
+
+
+def test_mistralai_embedding_documents_long() -> None:
+    """Test MistralAI embeddings for documents."""
+    documents = ["foo bar " * 1000, "test document " * 1000]
+    embedding = MistralAIEmbeddings()
+    output = embedding.embed_documents(documents)
+    assert len(output) == 2
+    assert len(output[0]) == 1024
+
+
+def test_mistralai_embed_query_character() -> None:
+    """Test MistralAI embeddings for query."""
+    document = "😳"
+    embedding = MistralAIEmbeddings()
+    output = embedding.embed_query(document)
+    assert len(output) == 1024