Make pyarrow dependency optional for tests (#2733) (#2735)

(cherry picked from commit d692c70)

Co-authored-by: danigm <[email protected]>

test_elasticsearch/test_client/test_deprecated_options.py

@@ -21,6 +21,23 @@
 
 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:
@@ -129,15 +146,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 +159,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

@@ -20,6 +20,23 @@
 from elasticsearch import Elasticsearch
 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):
@@ -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,15 @@
 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
@@ -32,12 +38,7 @@
 
 from elasticsearch import Elasticsearch
 from elasticsearch.exceptions import SerializationError
-from elasticsearch.serializer import (
-    JSONSerializer,
-    OrjsonSerializer,
-    PyArrowSerializer,
-    TextSerializer,
-)
+from elasticsearch.serializer import JSONSerializer, OrjsonSerializer, TextSerializer
 
 requires_numpy_and_pandas = pytest.mark.skipif(
     np is None or pd is None, reason="Test requires numpy and pandas to be available"
@@ -163,6 +164,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]),