Как сделать пользовательские сравнения в pytest?

Question

Например, я хотел бы заявить, что два Pyspark DataFrame имеют одинаковые данные, однако просто с помощью == проверяет, что это один и тот же объект.В идеале я также хотел бы указать, имеет ли значение порядок или нет.

Я попытался написать функцию, которая вызывает AssertionError, но добавляет шум к выходу pytest, поскольку он показывает трассировкуиз этой функции.

Еще одна мысль, которая у меня была, - это смоделировать метод __eq__ DataFrames, но я не уверен, что это правильный путь.

Редактировать:

Я рассмотрел просто использование функции, которая возвращает true или false вместо оператора, однако, похоже, что она не работает с pytest_assertrepr_compare.Я недостаточно знаком с тем, как работает этот хук, поэтому возможно, что есть способ использовать его с функцией вместо оператора.

Answer 1

просто используйте pandas.Dataframe.equals метод https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.equals.html

Например,

assert df1.equals(df2)

assert можно использовать со всем, что возвращает логическое значение.Так что да, вы можете написать любую пользовательскую функцию сравнения для сравнения двух объектов.Пока пользовательская функция возвращает логическое значение.Однако в этом случае нет необходимости в пользовательской функции, поскольку pandas уже предоставляет одну

Answer 2

Вы можете использовать один из ловушек pytest, в частности, pytest_assertrepr_compare .Там вы можете определить, что вы хотите сравнить и как, также документы довольно хороши и с примерами.Удачи.:)

Answer 3

Мое текущее решение - использовать патч для переопределения метода __eq__ DataFrame.Вот пример с Pandas, так как его быстрее тестировать, идея должна применяться к любому объекту.

import pandas as pd
# use this import for python3
# from unittest.mock import patch
from mock import patch


def custom_df_compare(self, other):
    # Put logic for comparing df's here
    # Returning True for demonstration
    return True


@patch("pandas.DataFrame.__eq__", custom_df_compare)
def test_df_equal():
    df1 = pd.DataFrame(
        {"id": [1, 2, 3], "name": ["a", "b", "c"]}, columns=["id", "name"]
    )
    df2 = pd.DataFrame(
        {"id": [2, 3, 4], "name": ["b", "c", "d"]}, columns=["id", "name"]
    )

    assert df1 == df2

Еще не пробовал, но планирую добавить его в качестве прибора и использовать autouse дляиспользуйте его для всех тестов автоматически.

Чтобы элегантно обработать индикатор «порядок имеет значение», я играю с подходом, похожим на pytest.approx, который возвращает новый класс со своим собственным __eq__, например:

class SortedDF(object):
    "Indicates that the order of data matters when comparing to another df"

    def __init__(self, df):
        self.df = df

    def __eq__(self, other):
        # Put logic for comparing df's including order of data here
        # Returning True for demonstration purposes
        return True


def test_sorted_df():
    df1 = pd.DataFrame(
        {"id": [1, 2, 3], "name": ["a", "b", "c"]}, columns=["id", "name"]
    )
    df2 = pd.DataFrame(
        {"id": [2, 3, 4], "name": ["b", "c", "d"]}, columns=["id", "name"]
    )

    # Passes because SortedDF.__eq__ is used
    assert SortedDF(df1) == df2
    # Fails because df2's __eq__ method is used
    assert df2 == SortedDF(df2)

Незначительная проблема, которую я не смог решить, - это ошибка второго утверждения, assert df2 == SortedDF(df2).Этот порядок отлично работает с pytest.approx, но не здесь.Я попытался прочитать оператор ==, но не смог выяснить, как исправить второй случай.

Answer 4

Чтобы выполнить необработанное сравнение значений DataFrames (должен быть точный порядок), вы можете сделать что-то вроде этого:

import pandas as pd
from pyspark.sql import Row

df1 = spark.createDataFrame([Row(a=1, b=2, c=3), Row(a=1, b=3, c=3)])
df2 = spark.createDataFrame([Row(a=1, b=2, c=3), Row(a=1, b=3, c=3)])

pd.testing.assert_frame_equal(df1.toPandas(), df2.toPandas())

Если вы хотите указать по порядку, вы можете сделать некоторые преобразованияв панде DataFrame сначала сортировать по определенному столбцу, используя следующую функцию:

def assert_frame_equal_with_sort(results, expected, keycolumns):
  results = results.reindex(sorted(results.columns), axis=1)
  expected = expected.reindex(sorted(expected.columns), axis=1)

  results_sorted = results.sort_values(by=keycolumns).reset_index(drop=True)
  expected_sorted = expected.sort_values(by=keycolumns).reset_index(drop=True)

  pd.testing.assert_frame_equal(results_sorted, expected_sorted)


df1 = spark.createDataFrame([Row(a=1, b=2, c=3), Row(a=1, b=3, c=3)])
df2 = spark.createDataFrame([Row(a=1, b=3, c=3), Row(a=1, b=2, c=3)])

assert_frame_equal_with_sort(df1.toPandas(), df2.toPandas(), ['b'])