Make pyarrow dependency optional for tests

Author: danigm

test_elasticsearch/test_client/test_deprecated_options.py

@@ -22,6 +22,24 @@
 from elasticsearch import Elasticsearch, JsonSerializer
 
 
+EXPECTED_SERIALIZERS = {
+    "application/vnd.mapbox-vector-tile",
+    "application/x-ndjson",
+    "application/json",
+    "text/*",
+    "application/vnd.elasticsearch+json",
+    "application/vnd.elasticsearch+x-ndjson",
+}
+
+
+try:
+    import pyarrow as pa
+    EXPECTED_SERIALIZERS.add("application/vnd.apache.arrow.stream")
+except ImportError:
+    pa = None
+
+
+
 def test_sniff_on_connection_fail():
     with warnings.catch_warnings(record=True) as w:
         client = Elasticsearch("http://localhost:9200", sniff_on_connection_fail=True)
@@ -129,15 +147,7 @@ class CustomSerializer(JsonSerializer):
         client.transport.serializers.get_serializer("application/json"),
         CustomSerializer,
     )
-    assert set(client.transport.serializers.serializers.keys()) == {
-        "application/vnd.mapbox-vector-tile",
-        "application/x-ndjson",
-        "application/json",
-        "text/*",
-        "application/vnd.apache.arrow.stream",
-        "application/vnd.elasticsearch+json",
-        "application/vnd.elasticsearch+x-ndjson",
-    }
+    assert set(client.transport.serializers.serializers.keys()) == EXPECTED_SERIALIZERS
 
     client = Elasticsearch(
         "http://localhost:9200",
@@ -150,13 +160,5 @@ class CustomSerializer(JsonSerializer):
         client.transport.serializers.get_serializer("application/json"),
         CustomSerializer,
     )
-    assert set(client.transport.serializers.serializers.keys()) == {
-        "application/vnd.mapbox-vector-tile",
-        "application/x-ndjson",
-        "application/json",
-        "text/*",
-        "application/vnd.apache.arrow.stream",
-        "application/vnd.elasticsearch+json",
-        "application/vnd.elasticsearch+x-ndjson",
-        "application/cbor",
-    }
+    expected = EXPECTED_SERIALIZERS | {"application/cbor"}
+    assert set(client.transport.serializers.serializers.keys()) == expected

test_elasticsearch/test_client/test_serializers.py

@@ -21,6 +21,23 @@
 from test_elasticsearch.test_cases import DummyTransportTestCase
 
 
+EXPECTED_SERIALIZERS = {
+    "application/json",
+    "text/*",
+    "application/x-ndjson",
+    "application/vnd.mapbox-vector-tile",
+    "application/vnd.elasticsearch+json",
+    "application/vnd.elasticsearch+x-ndjson",
+}
+
+
+try:
+    import pyarrow as pa
+    EXPECTED_SERIALIZERS.add("application/vnd.apache.arrow.stream")
+except ImportError:
+    pa = None
+
+
 class TestSerializers(DummyTransportTestCase):
     def test_compat_mode_on_by_default(self):
         calls = self.client.transport.calls
@@ -90,16 +107,8 @@ class CustomSerializer:
             "https://localhost:9200", serializers={f"application/{mime_subtype}": ser}
         )
         serializers = client.transport.serializers.serializers
-        assert set(serializers.keys()) == {
-            "application/json",
-            "text/*",
-            "application/x-ndjson",
-            "application/vnd.apache.arrow.stream",
-            "application/vnd.mapbox-vector-tile",
-            "application/vnd.elasticsearch+json",
-            "application/vnd.elasticsearch+x-ndjson",
-        }
 
+        assert set(serializers.keys()) == EXPECTED_SERIALIZERS
         assert serializers[f"application/{mime_subtype}"] is ser
         assert serializers[f"application/vnd.elasticsearch+{mime_subtype}"] is ser
 
@@ -118,16 +127,7 @@ class CustomSerializer:
             },
         )
         serializers = client.transport.serializers.serializers
-        assert set(serializers.keys()) == {
-            "application/json",
-            "text/*",
-            "application/x-ndjson",
-            "application/vnd.apache.arrow.stream",
-            "application/vnd.mapbox-vector-tile",
-            "application/vnd.elasticsearch+json",
-            "application/vnd.elasticsearch+x-ndjson",
-        }
-
+        assert set(serializers.keys()) == EXPECTED_SERIALIZERS
         assert serializers[f"application/{mime_subtype}"] is ser1
         assert serializers[f"application/vnd.elasticsearch+{mime_subtype}"] is ser2
 
@@ -138,15 +138,6 @@ class CustomSerializer:
         ser = CustomSerializer()
         client = Elasticsearch("https://localhost:9200", serializer=ser)
         serializers = client.transport.serializers.serializers
-        assert set(serializers.keys()) == {
-            "application/json",
-            "text/*",
-            "application/x-ndjson",
-            "application/vnd.apache.arrow.stream",
-            "application/vnd.mapbox-vector-tile",
-            "application/vnd.elasticsearch+json",
-            "application/vnd.elasticsearch+x-ndjson",
-        }
-
+        assert set(serializers.keys()) == EXPECTED_SERIALIZERS
         assert serializers["application/json"] is ser
         assert serializers["application/vnd.elasticsearch+json"] is ser

test_elasticsearch/test_serializer.py

@@ -19,9 +19,14 @@
 from datetime import datetime
 from decimal import Decimal
 
-import pyarrow as pa
 import pytest
 
+try:
+    import pyarrow as pa
+    from elasticsearch.serializer import PyArrowSerializer
+except ImportError:
+    pa = None
+
 try:
     import numpy as np
     import pandas as pd
@@ -35,7 +40,6 @@
 from elasticsearch.serializer import (
     JSONSerializer,
     OrjsonSerializer,
-    PyArrowSerializer,
     TextSerializer,
 )
 
@@ -163,6 +167,7 @@ def test_serializes_pandas_category(json_serializer):
     assert b'{"d":[1,2,3]}' == json_serializer.dumps({"d": cat})
 
 
+@pytest.mark.skipif(pa is None, reason="Test requires pyarrow to be available")
 def test_pyarrow_loads():
     data = [
         pa.array([1, 2, 3, 4]),