integration test

Author: miguelgrinberg

test_elasticsearch/test_dsl/test_integration/_async/test_document.py

@@ -25,6 +25,7 @@
 from ipaddress import ip_address
 from typing import TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Tuple, Union
 
+import numpy as np
 import pytest
 from pytest import raises
 from pytz import timezone
@@ -865,6 +866,7 @@ class Doc(AsyncDocument):
         float_vector: List[float] = mapped_field(DenseVector())
         byte_vector: List[int] = mapped_field(DenseVector(element_type="byte"))
         bit_vector: List[int] = mapped_field(DenseVector(element_type="bit"))
+        numpy_float_vector: np.ndarray = mapped_field(DenseVector(use_numpy=True))
 
         class Index:
             name = "vectors"
@@ -876,6 +878,7 @@ class Index:
         float_vector=[1.0, 1.2, 2.3],
         byte_vector=[12, 23, 34, 45],
         bit_vector=[18, -43, -112],
+        numpy_float_vector=np.array([3.1, 2.25, 1.0]),
     )
     await doc.save(refresh=True)
 
@@ -884,6 +887,8 @@ class Index:
     assert [round(v, 1) for v in docs[0].float_vector] == doc.float_vector
     assert docs[0].byte_vector == doc.byte_vector
     assert docs[0].bit_vector == doc.bit_vector
+    assert type(docs[0].numpy_float_vector) is np.ndarray
+    assert np.array_equal(docs[0].numpy_float_vector, doc.numpy_float_vector)
 
 
 @pytest.mark.anyio

test_elasticsearch/test_dsl/test_integration/_sync/test_document.py

@@ -25,6 +25,7 @@
 from ipaddress import ip_address
 from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union
 
+import numpy as np
 import pytest
 from pytest import raises
 from pytz import timezone
@@ -853,6 +854,7 @@ class Doc(Document):
         float_vector: List[float] = mapped_field(DenseVector())
         byte_vector: List[int] = mapped_field(DenseVector(element_type="byte"))
         bit_vector: List[int] = mapped_field(DenseVector(element_type="bit"))
+        numpy_float_vector: np.ndarray = mapped_field(DenseVector(use_numpy=True))
 
         class Index:
             name = "vectors"
@@ -864,6 +866,7 @@ class Index:
         float_vector=[1.0, 1.2, 2.3],
         byte_vector=[12, 23, 34, 45],
         bit_vector=[18, -43, -112],
+        numpy_float_vector=np.array([3.1, 2.25, 1.0]),
     )
     doc.save(refresh=True)
 
@@ -872,6 +875,8 @@ class Index:
     assert [round(v, 1) for v in docs[0].float_vector] == doc.float_vector
     assert docs[0].byte_vector == doc.byte_vector
     assert docs[0].bit_vector == doc.bit_vector
+    assert type(docs[0].numpy_float_vector) is np.ndarray
+    assert np.array_equal(docs[0].numpy_float_vector, doc.numpy_float_vector)
 
 
 @pytest.mark.sync