From dea2d04565abab2dee8944c3b2f332dd758681fe Mon Sep 17 00:00:00 2001 From: DShi Date: Mon, 25 Nov 2024 23:39:34 -0500 Subject: [PATCH] introduce MockRow class for object comparison checks --- .../tests/snowflake/operators/test_snowflake_sql.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/providers/tests/snowflake/operators/test_snowflake_sql.py b/providers/tests/snowflake/operators/test_snowflake_sql.py index 11f1859a75fbe..5deb39e541d1a 100644 --- a/providers/tests/snowflake/operators/test_snowflake_sql.py +++ b/providers/tests/snowflake/operators/test_snowflake_sql.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -from typing import Any from unittest import mock from unittest.mock import MagicMock, patch @@ -26,8 +25,15 @@ from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator -def MockRow(*args: Any, **kwargs: Any) -> MagicMock: - return MagicMock() +class MockRow: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + def __eq__(self, other): + return isinstance(other, MockRow) and self.__dict__ == other.__dict__ + + def __repr__(self): + return f"MockRow({self.__dict__})" from airflow.models.connection import Connection