From 3b4371aaddad13cb1c097027c6cff934b6b1fe5a Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 30 Apr 2023 19:01:15 +0100
Subject: [PATCH 01/10] Start making unit testing more general

---
 xarray/tests/test_units.py | 147 ++++++++++++++++++++++++++++++++-----
 1 file changed, 128 insertions(+), 19 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index ad9da0fad33..e1fe5df27ff 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -2,6 +2,9 @@
 
 import functools
 import operator
+from abc import ABC, abstractclassmethod
+from dataclasses import dataclass
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -31,6 +34,28 @@
 DimensionalityError = pint.errors.DimensionalityError
 
 
+@dataclass
+class UnitInfo(ABC):
+    unit: Any
+    unit_type: type
+    quantity_type: type
+
+    @abstractclassmethod
+    def strip_units(quantity):
+        # Remove units from a quantity
+        pass
+
+    @abstractclassmethod
+    def get_unit(quantity):
+        # Get the unit of a given quantity
+        pass
+
+    @abstractclassmethod
+    def assert_equal(q1, q2):
+        # Assert that two quantities are equal
+        pass
+
+
 # make sure scalars are converted to 0d arrays so quantities can
 # always be treated like ndarrays
 unit_registry = pint.UnitRegistry(force_ndarray_like=True)
@@ -42,6 +67,53 @@
 ]
 
 
+class PintInfo(UnitInfo):
+    unit = unit_registry.m
+    unit_type = pint.Unit
+    quantity_type = pint.Quantity
+
+    @staticmethod
+    def strip_units(quantity):
+        return quantity.magnitude
+
+    @staticmethod
+    def get_unit(quantity):
+        return quantity.units
+
+    @staticmethod
+    def assert_equal(q1, q2):
+        assert (q1 == q2).all()
+
+
+"""
+class AstropyInfo(UnitInfo):
+    unit = astropy.units.m
+    unit_type = astropy.units.UnitBase
+    quantity_type = astropy.units.Quantity
+
+    def strip_units(quantity):
+        return quantity.value
+
+    @staticmethod
+    def get_unit(quantity):
+        return quantity.unit
+
+    def assert_equal(q1, q2):
+        astropy.units.assert_quantity_equal(q1, q2)
+"""
+
+unit_libs = [PintInfo]  # + [AstropyInfo]
+known_quantities = tuple(lib.quantity_type for lib in unit_libs)
+
+
+@pytest.fixture(params=unit_libs)
+def unit_lib(request):
+    """
+    A fixture to return a unit represented in different libraries.
+    """
+    return request.param
+
+
 def is_compatible(unit1, unit2):
     def dimensionality(obj):
         if isinstance(obj, (unit_registry.Quantity, unit_registry.Unit)):
@@ -78,17 +150,19 @@ def array_extract_units(obj):
     if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
         obj = obj.data
 
-    try:
-        return obj.units
-    except AttributeError:
-        return None
+    for unit_lib in unit_libs:
+        if isinstance(obj, unit_lib.quantity_type):
+            return unit_lib.get_unit(obj)
+
+    return None
 
 
 def array_strip_units(array):
-    try:
-        return array.magnitude
-    except AttributeError:
-        return array
+    for unit_lib in unit_libs:
+        if isinstance(array, unit_lib.quantity_type):
+            return unit_lib.strip_units(array)
+
+    return array
 
 
 def array_attach_units(data, unit):
@@ -125,18 +199,24 @@ def extract_units(obj):
         units = {**vars_units, **coords_units}
     elif isinstance(obj, xr.Variable):
         vars_units = {None: array_extract_units(obj.data)}
-
-        units = {**vars_units}
-    elif isinstance(obj, Quantity):
-        vars_units = {None: array_extract_units(obj)}
-
         units = {**vars_units}
     else:
         units = {}
+        for unit_lib in unit_libs:
+            if isinstance(obj, unit_lib.quantity_type):
+                vars_units = {None: array_extract_units(obj)}
+                units = {**vars_units}
+                break
 
     return units
 
 
+def test_extract_units(unit_lib):
+    unit = unit_lib.unit
+    quantity = np.arange(10) * unit
+    assert extract_units(quantity) == {None: unit}
+
+
 def strip_units(obj):
     if isinstance(obj, xr.Dataset):
         data_vars = {
@@ -165,8 +245,8 @@ def strip_units(obj):
     elif isinstance(obj, xr.Variable):
         data = array_strip_units(obj.data)
         new_obj = obj.copy(data=data)
-    elif isinstance(obj, unit_registry.Quantity):
-        new_obj = obj.magnitude
+    elif isinstance(obj, known_quantities):
+        new_obj = array_strip_units(obj)
     elif isinstance(obj, (list, tuple)):
         return type(obj)(strip_units(elem) for elem in obj)
     else:
@@ -175,6 +255,19 @@ def strip_units(obj):
     return new_obj
 
 
+def test_strip_units(unit_lib):
+    unit = unit_lib.unit
+    # Array
+    quantity = np.arange(10) * unit
+    np.testing.assert_equal(strip_units(quantity), np.arange(10))
+
+    # DataArray
+    array = np.linspace(0, 10, 20) * unit
+    x = np.arange(20)
+    da = xr.DataArray(data=array, dims="x", coords={"x": x})
+    assert type(strip_units(da).data) == np.ndarray
+
+
 def attach_units(obj, units):
     if not isinstance(obj, (xr.DataArray, xr.Dataset, xr.Variable)):
         units = units.get("data", None) or units.get(None, None) or 1
@@ -219,6 +312,21 @@ def attach_units(obj, units):
     return new_obj
 
 
+def test_attach_units(unit_lib):
+    unit = unit = unit_lib.unit
+    # Array
+    array = np.arange(10)
+    quantity = attach_units(array, units={"data": unit})
+    unit_lib.assert_equal(quantity, np.arange(10) * unit)
+
+    # DataArray
+    array = np.linspace(0, 10, 20)
+    x = np.arange(20)
+    da = xr.DataArray(data=array, dims="x", coords={"x": x})
+    da = attach_units(da, {None: unit})
+    assert isinstance(da.data, unit_lib.quantity_type)
+
+
 def convert_units(obj, to):
     # preprocess
     to = {
@@ -382,11 +490,12 @@ def __repr__(self):
         "coords",
     ),
 )
-def test_apply_ufunc_dataarray(variant, dtype):
+def test_apply_ufunc_dataarray(variant, unit_lib, dtype):
+    unit = unit_lib.unit
     variants = {
-        "data": (unit_registry.m, 1, 1),
-        "dims": (1, unit_registry.m, 1),
-        "coords": (1, 1, unit_registry.m),
+        "data": (unit, 1, 1),
+        "dims": (1, unit, 1),
+        "coords": (1, 1, unit),
     }
     data_unit, dim_unit, coord_unit = variants.get(variant)
     func = functools.partial(

From 013debf13f29708bf31db9f6868be0aea55f367c Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Mon, 1 May 2023 17:24:39 +0100
Subject: [PATCH 02/10] More unit test generalisation

---
 xarray/tests/test_units.py | 74 ++++++++++++++++++++++++--------------
 1 file changed, 48 insertions(+), 26 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index e1fe5df27ff..a0ea2338d30 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -36,7 +36,14 @@
 
 @dataclass
 class UnitInfo(ABC):
+    # A unit (e.g. m)
     unit: Any
+    # A scaled version of unit (e.g. mm)
+    scaled_unit: Any
+    # A unit that can't be converted to unit (e.g. s)
+    incompatible_unit: Any
+    # Dimensionless unit
+    dimensionless: Any
     unit_type: type
     quantity_type: type
 
@@ -69,6 +76,10 @@ def assert_equal(q1, q2):
 
 class PintInfo(UnitInfo):
     unit = unit_registry.m
+    scaled_unit = unit_registry.mm
+    incompatible_unit = unit_registry.s
+    dimensionless = unit_registry.dimensionless
+
     unit_type = pint.Unit
     quantity_type = pint.Quantity
 
@@ -88,6 +99,10 @@ def assert_equal(q1, q2):
 """
 class AstropyInfo(UnitInfo):
     unit = astropy.units.m
+    scaled_unit = astropy.units.mm
+    incompatible_unit = astropy.units.s
+    dimensionless = astropy.units.dimensionless_unscaled
+
     unit_type = astropy.units.UnitBase
     quantity_type = astropy.units.Quantity
 
@@ -98,12 +113,14 @@ def strip_units(quantity):
     def get_unit(quantity):
         return quantity.unit
 
+    @staticmethod
     def assert_equal(q1, q2):
-        astropy.units.assert_quantity_equal(q1, q2)
+        (q1 == q2).all()
 """
 
 unit_libs = [PintInfo]  # + [AstropyInfo]
-known_quantities = tuple(lib.quantity_type for lib in unit_libs)
+known_quantity_types = tuple(lib.quantity_type for lib in unit_libs)
+known_unit_types = tuple(lib.unit_type for lib in unit_libs)
 
 
 @pytest.fixture(params=unit_libs)
@@ -233,7 +250,7 @@ def strip_units(obj):
         coords = {
             strip_units(name): (
                 (value.dims, array_strip_units(value.variable._data))
-                if isinstance(value.data, Quantity)
+                if isinstance(value.data, known_quantity_types)
                 else value  # to preserve multiindexes
             )
             for name, value in obj.coords.items()
@@ -245,7 +262,7 @@ def strip_units(obj):
     elif isinstance(obj, xr.Variable):
         data = array_strip_units(obj.data)
         new_obj = obj.copy(data=data)
-    elif isinstance(obj, known_quantities):
+    elif isinstance(obj, known_quantity_types):
         new_obj = array_strip_units(obj)
     elif isinstance(obj, (list, tuple)):
         return type(obj)(strip_units(elem) for elem in obj)
@@ -524,11 +541,12 @@ def test_apply_ufunc_dataarray(variant, unit_lib, dtype):
         "coords",
     ),
 )
-def test_apply_ufunc_dataset(variant, dtype):
+def test_apply_ufunc_dataset(variant, unit_lib, dtype):
+    unit = unit_lib.unit
     variants = {
-        "data": (unit_registry.m, 1, 1),
-        "dims": (1, unit_registry.m, 1),
-        "coords": (1, 1, unit_registry.s),
+        "data": (unit, 1, 1),
+        "dims": (1, unit, 1),
+        "coords": (1, 1, unit),
     }
     data_unit, dim_unit, coord_unit = variants.get(variant)
 
@@ -560,12 +578,10 @@ def test_apply_ufunc_dataset(variant, dtype):
     "unit,error",
     (
         pytest.param(1, DimensionalityError, id="no_unit"),
-        pytest.param(
-            unit_registry.dimensionless, DimensionalityError, id="dimensionless"
-        ),
-        pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
-        pytest.param(unit_registry.mm, None, id="compatible_unit"),
-        pytest.param(unit_registry.m, None, id="identical_unit"),
+        pytest.param("dimensionless", DimensionalityError, id="dimensionless"),
+        pytest.param("incompatible_unit", DimensionalityError, id="incompatible_unit"),
+        pytest.param("scaled_unit", None, id="compatible_unit"),
+        pytest.param("unit", None, id="identical_unit"),
     ),
     ids=repr,
 )
@@ -580,9 +596,12 @@ def test_apply_ufunc_dataset(variant, dtype):
     ),
 )
 @pytest.mark.parametrize("value", (10, dtypes.NA))
-def test_align_dataarray(value, variant, unit, error, dtype):
+def test_align_dataarray(value, variant, unit, error, dtype, unit_lib):
+    if isinstance(unit, str):
+        unit = getattr(unit_lib, unit)
+
     if variant == "coords" and (
-        value != dtypes.NA or isinstance(unit, unit_registry.Unit)
+        value != dtypes.NA or isinstance(unit, known_unit_types)
     ):
         pytest.xfail(
             reason=(
@@ -593,7 +612,7 @@ def test_align_dataarray(value, variant, unit, error, dtype):
 
     fill_value = dtypes.get_fill_value(dtype) if value == dtypes.NA else value
 
-    original_unit = unit_registry.m
+    original_unit = unit_lib.unit
 
     variants = {
         "data": ((original_unit, unit), (1, 1), (1, 1)),
@@ -686,7 +705,7 @@ def test_align_dataarray(value, variant, unit, error, dtype):
 @pytest.mark.parametrize("value", (10, dtypes.NA))
 def test_align_dataset(value, unit, variant, error, dtype):
     if variant == "coords" and (
-        value != dtypes.NA or isinstance(unit, unit_registry.Unit)
+        value != dtypes.NA or isinstance(unit, known_unit_types)
     ):
         pytest.xfail(
             reason=(
@@ -764,10 +783,11 @@ def test_align_dataset(value, unit, variant, error, dtype):
     assert_allclose(expected_b, actual_b)
 
 
-def test_broadcast_dataarray(dtype):
+def test_broadcast_dataarray(dtype, unit_lib):
+    unit = unit_lib.unit
     # uses align internally so more thorough tests are not needed
-    array1 = np.linspace(0, 10, 2) * unit_registry.Pa
-    array2 = np.linspace(0, 10, 3) * unit_registry.Pa
+    array1 = np.linspace(0, 10, 2) * unit
+    array2 = np.linspace(0, 10, 3) * unit
 
     a = xr.DataArray(data=array1, dims="x")
     b = xr.DataArray(data=array2, dims="y")
@@ -786,10 +806,12 @@ def test_broadcast_dataarray(dtype):
     assert_identical(expected_b, actual_b)
 
 
-def test_broadcast_dataset(dtype):
+def test_broadcast_dataset(dtype, unit_lib):
+    unit = unit_lib.unit
+    scaled_unit = unit_lib.scaled_unit
     # uses align internally so more thorough tests are not needed
-    array1 = np.linspace(0, 10, 2) * unit_registry.Pa
-    array2 = np.linspace(0, 10, 3) * unit_registry.Pa
+    array1 = np.linspace(0, 10, 2) * unit
+    array2 = np.linspace(0, 10, 3) * unit
 
     x1 = np.arange(2)
     y1 = np.arange(3)
@@ -802,8 +824,8 @@ def test_broadcast_dataset(dtype):
     )
     other = xr.Dataset(
         data_vars={
-            "a": ("x", array1.to(unit_registry.hPa)),
-            "b": ("y", array2.to(unit_registry.hPa)),
+            "a": ("x", array1.to(scaled_unit)),
+            "b": ("y", array2.to(scaled_unit)),
         },
         coords={"x": x2, "y": y2},
     )

From 11988c208913e3cf3cee3feb00e122e64b13fd78 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 14 May 2023 11:16:23 +0100
Subject: [PATCH 03/10] .all() > np.all()

---
 xarray/tests/test_units.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index a0ea2338d30..96a00a9292b 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -93,7 +93,7 @@ def get_unit(quantity):
 
     @staticmethod
     def assert_equal(q1, q2):
-        assert (q1 == q2).all()
+        assert np.all(q1 == q2)
 
 
 """

From ff1e4db0d431cf8570d819ba777b3fc16e8fa332 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 14 May 2023 11:18:18 +0100
Subject: [PATCH 04/10] scaled_unit > compatible_unit

---
 xarray/tests/test_units.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 96a00a9292b..74f0ac1568d 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -38,8 +38,8 @@
 class UnitInfo(ABC):
     # A unit (e.g. m)
     unit: Any
-    # A scaled version of unit (e.g. mm)
-    scaled_unit: Any
+    # A different unit that previous unit can be transformed to (e.g. mm)
+    compatible_unit: Any
     # A unit that can't be converted to unit (e.g. s)
     incompatible_unit: Any
     # Dimensionless unit
@@ -76,7 +76,7 @@ def assert_equal(q1, q2):
 
 class PintInfo(UnitInfo):
     unit = unit_registry.m
-    scaled_unit = unit_registry.mm
+    compatible_unit = unit_registry.mm
     incompatible_unit = unit_registry.s
     dimensionless = unit_registry.dimensionless
 
@@ -99,7 +99,7 @@ def assert_equal(q1, q2):
 """
 class AstropyInfo(UnitInfo):
     unit = astropy.units.m
-    scaled_unit = astropy.units.mm
+    compatible_unit = astropy.units.mm
     incompatible_unit = astropy.units.s
     dimensionless = astropy.units.dimensionless_unscaled
 
@@ -580,7 +580,7 @@ def test_apply_ufunc_dataset(variant, unit_lib, dtype):
         pytest.param(1, DimensionalityError, id="no_unit"),
         pytest.param("dimensionless", DimensionalityError, id="dimensionless"),
         pytest.param("incompatible_unit", DimensionalityError, id="incompatible_unit"),
-        pytest.param("scaled_unit", None, id="compatible_unit"),
+        pytest.param("compatible_unit", None, id="compatible_unit"),
         pytest.param("unit", None, id="identical_unit"),
     ),
     ids=repr,
@@ -808,7 +808,7 @@ def test_broadcast_dataarray(dtype, unit_lib):
 
 def test_broadcast_dataset(dtype, unit_lib):
     unit = unit_lib.unit
-    scaled_unit = unit_lib.scaled_unit
+    compatible_unit = unit_lib.compatible_unit
     # uses align internally so more thorough tests are not needed
     array1 = np.linspace(0, 10, 2) * unit
     array2 = np.linspace(0, 10, 3) * unit
@@ -824,8 +824,8 @@ def test_broadcast_dataset(dtype, unit_lib):
     )
     other = xr.Dataset(
         data_vars={
-            "a": ("x", array1.to(scaled_unit)),
-            "b": ("y", array2.to(scaled_unit)),
+            "a": ("x", array1.to(compatible_unit)),
+            "b": ("y", array2.to(compatible_unit)),
         },
         coords={"x": x2, "y": y2},
     )
@@ -4109,21 +4109,21 @@ class TestDataset:
     )
     def test_init(self, shared, unit, error, dtype):
         original_unit = unit_registry.m
-        scaled_unit = unit_registry.mm
+        compatible_unit = unit_registry.mm
 
         a = np.linspace(0, 1, 10).astype(dtype) * unit_registry.Pa
         b = np.linspace(-1, 0, 10).astype(dtype) * unit_registry.degK
 
         values_a = np.arange(a.shape[0])
         dim_a = values_a * original_unit
-        coord_a = dim_a.to(scaled_unit)
+        coord_a = dim_a.to(compatible_unit)
 
         values_b = np.arange(b.shape[0])
         dim_b = values_b * unit
         coord_b = (
-            dim_b.to(scaled_unit)
-            if unit_registry.is_compatible_with(dim_b, scaled_unit)
-            and unit != scaled_unit
+            dim_b.to(compatible_unit)
+            if unit_registry.is_compatible_with(dim_b, compatible_unit)
+            and unit != compatible_unit
             else dim_b * 1000
         )
 

From 7d2c211333a359aa66f10b4f599e0205e24e664e Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 14 May 2023 11:26:07 +0100
Subject: [PATCH 05/10] Remove some for loops

---
 xarray/tests/test_units.py | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 74f0ac1568d..5b094606b49 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -122,6 +122,13 @@ def assert_equal(q1, q2):
 known_quantity_types = tuple(lib.quantity_type for lib in unit_libs)
 known_unit_types = tuple(lib.unit_type for lib in unit_libs)
 
+def get_unit_lib(obj):
+    for unit_lib in unit_libs:
+        if isinstance(obj, unit_lib.quantity_type):
+            return unit_lib
+
+    return None
+
 
 @pytest.fixture(params=unit_libs)
 def unit_lib(request):
@@ -167,19 +174,19 @@ def array_extract_units(obj):
     if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
         obj = obj.data
 
-    for unit_lib in unit_libs:
-        if isinstance(obj, unit_lib.quantity_type):
-            return unit_lib.get_unit(obj)
-
-    return None
+    lib = get_unit_lib(obj)
+    if lib is not None:
+        return lib.get_unit(obj)
+    else:
+        return None
 
 
 def array_strip_units(array):
-    for unit_lib in unit_libs:
-        if isinstance(array, unit_lib.quantity_type):
-            return unit_lib.strip_units(array)
-
-    return array
+    lib = get_unit_lib(array)
+    if lib is not None:
+        return lib.strip_units(array)
+    else:
+        return array
 
 
 def array_attach_units(data, unit):

From 7ba75da379059e4654d9fef55fdd78a73bc66cc8 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sun, 14 May 2023 10:29:06 +0000
Subject: [PATCH 06/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 xarray/tests/test_units.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 5b094606b49..267624b1cdd 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -122,6 +122,7 @@ def assert_equal(q1, q2):
 known_quantity_types = tuple(lib.quantity_type for lib in unit_libs)
 known_unit_types = tuple(lib.unit_type for lib in unit_libs)
 
+
 def get_unit_lib(obj):
     for unit_lib in unit_libs:
         if isinstance(obj, unit_lib.quantity_type):

From e6ea73efe87f52d51feb61a960bf159fb70a1f05 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Sun, 14 May 2023 11:43:08 +0100
Subject: [PATCH 07/10] Fix assert_allclose calls

---
 xarray/tests/test_units.py | 32 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 267624b1cdd..ee8d9604a69 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -682,9 +682,9 @@ def test_align_dataarray(value, variant, unit, error, dtype, unit_lib):
     actual_a, actual_b = func(data_array1, data_array2)
 
     assert_units_equal(expected_a, actual_a)
-    assert_allclose(expected_a, actual_a)
+    assert_allclose(expected_a, actual_a, atol=0)
     assert_units_equal(expected_b, actual_b)
-    assert_allclose(expected_b, actual_b)
+    assert_allclose(expected_b, actual_b, atol=0)
 
 
 @pytest.mark.parametrize(
@@ -786,9 +786,9 @@ def test_align_dataset(value, unit, variant, error, dtype):
     actual_a, actual_b = func(ds1, ds2)
 
     assert_units_equal(expected_a, actual_a)
-    assert_allclose(expected_a, actual_a)
+    assert_allclose(expected_a, actual_a, atol=0)
     assert_units_equal(expected_b, actual_b)
-    assert_allclose(expected_b, actual_b)
+    assert_allclose(expected_b, actual_b, atol=0)
 
 
 def test_broadcast_dataarray(dtype, unit_lib):
@@ -1270,7 +1270,7 @@ def test_merge_dataarray(variant, unit, error, dtype):
     actual = xr.merge([arr1, arr2, arr3])
 
     assert_units_equal(expected, actual)
-    assert_allclose(expected, actual)
+    assert_allclose(expected, actual, atol=0)
 
 
 @pytest.mark.parametrize(
@@ -1360,7 +1360,7 @@ def test_merge_dataset(variant, unit, error, dtype):
     actual = func([ds1, ds2, ds3])
 
     assert_units_equal(expected, actual)
-    assert_allclose(expected, actual)
+    assert_allclose(expected, actual, atol=0)
 
 
 @pytest.mark.parametrize(
@@ -1690,7 +1690,7 @@ def test_aggregation(self, func, dtype):
         actual = func(variable)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     def test_aggregate_complex(self):
         variable = xr.Variable("x", [1, 2j, np.nan] * unit_registry.m)
@@ -1698,7 +1698,7 @@ def test_aggregate_complex(self):
         actual = variable.mean()
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize(
         "func",
@@ -1756,7 +1756,7 @@ def test_numpy_methods(self, func, unit, error, dtype):
         actual = func(variable, *args, **kwargs)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize(
         "func", (method("item", 5), method("searchsorted", 5)), ids=repr
@@ -2079,7 +2079,7 @@ def test_1d_math(self, func, unit, error, dtype):
         actual = func(variable, y)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize(
         "unit,error",
@@ -2553,7 +2553,7 @@ def test_aggregation(self, func, dtype):
         actual = func(data_array)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize(
         "func",
@@ -3636,7 +3636,7 @@ def test_interp_reindex(self, variant, func, dtype):
         actual = func(data_array, x=new_x)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.skip(reason="indexes don't support units")
     @pytest.mark.parametrize(
@@ -3712,7 +3712,7 @@ def test_interp_reindex_like(self, variant, func, dtype):
         actual = func(data_array, other)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.skip(reason="indexes don't support units")
     @pytest.mark.parametrize(
@@ -4017,7 +4017,7 @@ def test_computation_objects(self, func, variant, dtype):
         actual = func(data_array).mean()
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     def test_resample(self, dtype):
         array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m
@@ -4270,7 +4270,7 @@ def test_aggregation(self, func, dtype):
         expected = attach_units(func(strip_units(ds)), units)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize("property", ("imag", "real"))
     def test_numpy_properties(self, property, dtype):
@@ -5505,7 +5505,7 @@ def test_computation_objects(self, func, variant, dtype):
         actual = func(ds).mean(*args, **kwargs)
 
         assert_units_equal(expected, actual)
-        assert_allclose(expected, actual)
+        assert_allclose(expected, actual, atol=0)
 
     @pytest.mark.parametrize(
         "variant",

From c80a98bf1a683e82e7e8031ac3497fd47171beda Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Wed, 28 Jun 2023 15:05:40 +0100
Subject: [PATCH 08/10] Fix pint subclassing

---
 xarray/tests/test_units.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index ee8d9604a69..5c2bacd5edf 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -80,8 +80,8 @@ class PintInfo(UnitInfo):
     incompatible_unit = unit_registry.s
     dimensionless = unit_registry.dimensionless
 
-    unit_type = pint.Unit
-    quantity_type = pint.Quantity
+    unit_type = unit_registry.Unit
+    quantity_type = unit_registry.Quantity
 
     @staticmethod
     def strip_units(quantity):

From cdf9cad204a9b67aaed16c73d3851f6c36d4ec9e Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Wed, 28 Jun 2023 15:09:03 +0100
Subject: [PATCH 09/10] Remove redundant loop

---
 xarray/tests/test_units.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 5c2bacd5edf..1440d6924f5 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -226,12 +226,8 @@ def extract_units(obj):
         vars_units = {None: array_extract_units(obj.data)}
         units = {**vars_units}
     else:
-        units = {}
-        for unit_lib in unit_libs:
-            if isinstance(obj, unit_lib.quantity_type):
-                vars_units = {None: array_extract_units(obj)}
-                units = {**vars_units}
-                break
+        vars_units = {None: array_extract_units(obj)}
+        units = {**vars_units}
 
     return units
 

From 038b3db3491723a9d4eb0c94a9808a1c812b7908 Mon Sep 17 00:00:00 2001
From: David Stansby <dstansby@gmail.com>
Date: Wed, 28 Jun 2023 15:15:33 +0100
Subject: [PATCH 10/10] Use alias for unit_registry

---
 xarray/tests/test_units.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 1440d6924f5..dda89a28f61 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -75,13 +75,15 @@ def assert_equal(q1, q2):
 
 
 class PintInfo(UnitInfo):
-    unit = unit_registry.m
-    compatible_unit = unit_registry.mm
-    incompatible_unit = unit_registry.s
-    dimensionless = unit_registry.dimensionless
+    ureg = unit_registry
+
+    unit = ureg.m
+    compatible_unit = ureg.mm
+    incompatible_unit = ureg.s
+    dimensionless = ureg.dimensionless
 
-    unit_type = unit_registry.Unit
-    quantity_type = unit_registry.Quantity
+    unit_type = ureg.Unit
+    quantity_type = ureg.Quantity
 
     @staticmethod
     def strip_units(quantity):