Introducing fixture to create `Connections` without DB in provider tests #51930

amoghrajesh · 2025-06-19T11:05:14Z

Why?

While working on moving basehook to task sdk (#51873), I realized that most providers set up their connections using db.merge_conn() with thousands of occurrences across the codebase.

db.merge_conn() is a utility that creates connections in the database for provider tests:

Overall slows down test execution due to database transactions
Can cause test isolation issues when connections persist between tests
Require database setup/teardown, making tests more complex
Can lead to flaky tests due to DB access issues

My initial idea was to create a fixture that could intercept, find, and replace those calls with task SDK comms calls automatically. However, I decided against this approach because it didn't add much value and was not maintainable.

How?

Created a new pytest fixture that leveraged the fact that Airflow has a Env backend setup by default with precedence being higher than that of the database: https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html#storing-connections-in-environment-variables.

The create_connection_without_db fixture sets AIRFLOW_CONN_{conn_id.upper()} to conn.get_uri() which can be utilised during lookup.

Benefits:

Connections are resolved from environment variables first
No database state is required for connection setup
Tests are more isolated and predictable

Effort involved:

There were various kinds of occurrences of db.merge_conn:

In unit tests:
- setup_class
- setup_method
- setup_module
- In local functions
- And some arbitrary ones
In system tests
In example dags!

TODO

Handle occurrences in unit tests
Have some code level TODOs, need to be checked in follow up
Handle occurrences In system tests
Consider improving the fixture to take a list of connections and add them at once.
Handle for telegram provider too: providers/telegram/tests/unit/telegram/hooks/test_telegram.py

Error:

_ TestTelegramHook.test_should_raise_exception_if_chat_id_is_not_provided_anywhere _
providers/telegram/tests/unit/telegram/hooks/test_telegram.py:93: in test_should_raise_exception_if_chat_id_is_not_provided_anywhere
    hook.send_message({"text": "test telegram message"})
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:338: in wrapped_f
    return copy(f, *args, **kw)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:477: in __call__
    do = self.iter(retry_state=retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:378: in iter
    result = action(retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:400: in <lambda>
    self._add_action_func(lambda rs: rs.outcome.result())
/usr/local/lib/python3.9/concurrent/futures/_base.py:439: in result
    return self.__get_result()
/usr/local/lib/python3.9/concurrent/futures/_base.py:391: in __get_result
    raise self._exception
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:480: in __call__
    result = fn(*args, **kwargs)
providers/telegram/src/airflow/providers/telegram/hooks/telegram.py:159: in send_message
    response = asyncio.run(self.connection.send_message(**kwargs))
/usr/local/lib/python3.9/asyncio/runners.py:37: in run
    raise ValueError("a coroutine was expected, got {!r}".format(main))
E   ValueError: a coroutine was expected, got <MagicMock name='get_conn().send_message()' id='140457991957568'>
----------------------------- Captured stdout call -----------------------------
[2025-06-20T11:32:10.235+0000] {base.py:65} INFO - Connection Retrieved 'telegram_default'
[2025-06-20T11:32:10.236+0000] {base.py:65} INFO - Connection Retrieved 'telegram_default'
------------------------------ Captured log call -------------------------------
INFO     airflow.hooks.base:base.py:65 Connection Retrieved 'telegram_default'
INFO     airflow.hooks.base:base.py:65 Connection Retrieved 'telegram_default'
_ TestTelegramHook.test_should_raise_exception_if_chat_id_is_not_provided_anywhere_when_sending_file _
providers/telegram/tests/unit/telegram/hooks/test_telegram.py:208: in test_should_raise_exception_if_chat_id_is_not_provided_anywhere_when_sending_file
    hook.send_file({"file": "/file/to/path"})
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:338: in wrapped_f
    return copy(f, *args, **kw)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:477: in __call__
    do = self.iter(retry_state=retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:378: in iter
    result = action(retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:400: in <lambda>
    self._add_action_func(lambda rs: rs.outcome.result())
/usr/local/lib/python3.9/concurrent/futures/_base.py:439: in result
    return self.__get_result()
/usr/local/lib/python3.9/concurrent/futures/_base.py:391: in __get_result
    raise self._exception
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:480: in __call__
    result = fn(*args, **kwargs)
providers/telegram/src/airflow/providers/telegram/hooks/telegram.py:189: in send_file
    response = asyncio.run(self.connection.send_document(**kwargs))
/usr/local/lib/python3.9/asyncio/runners.py:37: in run
    raise ValueError("a coroutine was expected, got {!r}".format(main))
E   ValueError: a coroutine was expected, got <MagicMock name='get_conn().send_document()' id='140457990581216'>

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

potiuk · 2025-06-19T11:12:00Z

nice :)

potiuk · 2025-06-19T11:12:47Z

Hopefully we will be able to move most of the tests out of the "Db" zone this way :)

amoghrajesh · 2025-06-19T11:13:11Z

WARNING: This one is going to get crazy with the number of changes. Planning to use env vars instead of DB.

potiuk · 2025-06-19T11:16:38Z

WARNING: This one is going to get crazy with the number of changes. Planning to use env vars instead of DB.

Yeah. It's going to be big. Maybe consider implementing a parallel solution and make an issue where others will convert provider by provider? That would be way more manageable (and if conversion is simple, would be rather fast).

amoghrajesh · 2025-06-19T13:06:25Z

WARNING: This one is going to get crazy with the number of changes. Planning to use env vars instead of DB.

Yeah. It's going to be big. Maybe consider implementing a parallel solution and make an issue where others will convert provider by provider? That would be way more manageable (and if conversion is simple, would be rather fast).

Let me create an MVP off of this and try to figure out if it needs that kind of splitting.

providers/asana/tests/unit/asana/operators/test_asana_tasks.py

amoghrajesh · 2025-06-20T13:29:59Z

Ok, getting there now. Probably really close this time. I am running into some errors with the telegram provider which I cannot seem to figure out.

_ TestTelegramHook.test_should_raise_exception_if_chat_id_is_not_provided_anywhere _
providers/telegram/tests/unit/telegram/hooks/test_telegram.py:93: in test_should_raise_exception_if_chat_id_is_not_provided_anywhere
    hook.send_message({"text": "test telegram message"})
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:338: in wrapped_f
    return copy(f, *args, **kw)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:477: in __call__
    do = self.iter(retry_state=retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:378: in iter
    result = action(retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:400: in <lambda>
    self._add_action_func(lambda rs: rs.outcome.result())
/usr/local/lib/python3.9/concurrent/futures/_base.py:439: in result
    return self.__get_result()
/usr/local/lib/python3.9/concurrent/futures/_base.py:391: in __get_result
    raise self._exception
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:480: in __call__
    result = fn(*args, **kwargs)
providers/telegram/src/airflow/providers/telegram/hooks/telegram.py:159: in send_message
    response = asyncio.run(self.connection.send_message(**kwargs))
/usr/local/lib/python3.9/asyncio/runners.py:37: in run
    raise ValueError("a coroutine was expected, got {!r}".format(main))
E   ValueError: a coroutine was expected, got <MagicMock name='get_conn().send_message()' id='140457991957568'>
----------------------------- Captured stdout call -----------------------------
[2025-06-20T11:32:10.235+0000] {base.py:65} INFO - Connection Retrieved 'telegram_default'
[2025-06-20T11:32:10.236+0000] {base.py:65} INFO - Connection Retrieved 'telegram_default'
------------------------------ Captured log call -------------------------------
INFO     airflow.hooks.base:base.py:65 Connection Retrieved 'telegram_default'
INFO     airflow.hooks.base:base.py:65 Connection Retrieved 'telegram_default'
_ TestTelegramHook.test_should_raise_exception_if_chat_id_is_not_provided_anywhere_when_sending_file _
providers/telegram/tests/unit/telegram/hooks/test_telegram.py:208: in test_should_raise_exception_if_chat_id_is_not_provided_anywhere_when_sending_file
    hook.send_file({"file": "/file/to/path"})
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:338: in wrapped_f
    return copy(f, *args, **kw)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:477: in __call__
    do = self.iter(retry_state=retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:378: in iter
    result = action(retry_state)
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:400: in <lambda>
    self._add_action_func(lambda rs: rs.outcome.result())
/usr/local/lib/python3.9/concurrent/futures/_base.py:439: in result
    return self.__get_result()
/usr/local/lib/python3.9/concurrent/futures/_base.py:391: in __get_result
    raise self._exception
/usr/local/lib/python3.9/site-packages/tenacity/__init__.py:480: in __call__
    result = fn(*args, **kwargs)
providers/telegram/src/airflow/providers/telegram/hooks/telegram.py:189: in send_file
    response = asyncio.run(self.connection.send_document(**kwargs))
/usr/local/lib/python3.9/asyncio/runners.py:37: in run
    raise ValueError("a coroutine was expected, got {!r}".format(main))
E   ValueError: a coroutine was expected, got <MagicMock name='get_conn().send_document()' id='140457990581216'>

Will add it part of TODO and handle there.

amoghrajesh · 2025-06-20T13:31:01Z

@potiuk I am close with this one, do you wanna take a look?

jason810496

LGTM if CI is green.

providers/apache/kafka/tests/integration/apache/kafka/operators/test_consume.py

amoghrajesh · 2025-06-20T14:25:05Z

Oh the compat tests! I made changes in #51953, which should fix this but I wonder how i can apply those here.

amoghrajesh · 2025-06-20T17:05:15Z

OK, I decided to modify the fixture to use: as_json because it’s independent of the core changes in #51953 and generally quite stable. The only concern I mentioned is in this comment, but I don’t really have much of a choice at this point to avoid it from happening, it ends up being a bit of a chicken and egg situation here. Core isn't released but provider wants to use it.

Maybe if we see issues in future we can get back to uri repr.

CC @ashb

amoghrajesh · 2025-06-20T18:04:26Z

Finally green!

potiuk

NIIIIICE!

potiuk · 2025-06-21T18:33:15Z

Shall we merge ;) ?

amoghrajesh · 2025-06-22T08:32:35Z

Merging it, thanks for the reviews @potiuk @ashb @jason810496!

Yep @potiuk on it!

amoghrajesh · 2025-06-22T08:33:09Z

Will work on the follow up items starting next week

potiuk · 2025-06-22T09:55:13Z

Will work on the follow up items starting next week

Let me see how much od the providers tests I can now move to be non-db tests :D.... That's on me

After apache#51930, we can now remove "pytest.mark.db_test" in all providers that used db only to create connections. This PR is the "trial" attempt of doing so and it verifies that it is a viable thing to do. It: * removes all pytest.mark.db_test markers from airbyte * removes / reshuffles many pytest.mark.db_tests in amazon * adds pre-commit that we can use to guard that no new pytest.mark.db_test markers will be added in providers that we already cleaned up (for example airbyte) This PR, when merged will be followed up with an issue where we will ask contributors to apply the same approach to all the remaining providers - where I will describe in detail the process of removing the markers. This is part of the apache#42632 which has the long-term target of making all the provider tests non-db tests and simplifying our test setup.

potiuk · 2025-06-22T11:07:11Z

Will work on the follow up items starting next week

Let me see how much od the providers tests I can now move to be non-db tests :D.... That's on me

First phase - amazon + airbyte here #52017

After apache#51930, we can now remove "pytest.mark.db_test" in all providers that used db only to create connections. This PR is the "trial" attempt of doing so and it verifies that it is a viable thing to do. It: * removes all pytest.mark.db_test markers from airbyte * removes / reshuffles many pytest.mark.db_tests in amazon * adds pre-commit that we can use to guard that no new pytest.mark.db_test markers will be added in providers that we already cleaned up (for example airbyte) This PR, when merged will be followed up with an issue where we will ask contributors to apply the same approach to all the remaining providers - where I will describe in detail the process of removing the markers. This is part of the apache#42632 which has the long-term target of making all the provider tests non-db tests and simplifying our test setup.

…le (#52017) After #51930, we can now remove "pytest.mark.db_test" in all providers that used db only to create connections. This PR is the "trial" attempt of doing so and it verifies that it is a viable thing to do. It: * removes all pytest.mark.db_test markers from airbyte * removes / reshuffles many pytest.mark.db_tests in amazon * adds pre-commit that we can use to guard that no new pytest.mark.db_test markers will be added in providers that we already cleaned up (for example airbyte) This PR, when merged will be followed up with an issue where we will ask contributors to apply the same approach to all the remaining providers - where I will describe in detail the process of removing the markers. This is part of the #42632 which has the long-term target of making all the provider tests non-db tests and simplifying our test setup.

amoghrajesh · 2025-06-23T08:37:08Z

@potiuk do you think we have to make changes to the system tests too?
I think those are using the DB and those seem to be ok?

potiuk · 2025-06-23T08:41:13Z

@potiuk do you think we have to make changes to the system tests too? I think those are using the DB and those seem to be ok?

Well. Not really - those are regular DAGs, so they generally should not need (and cannot) access the DB. So we should adapt them too.

amoghrajesh · 2025-06-23T08:43:19Z

On it!

amoghrajesh added 2 commits June 19, 2025 16:34

Adding a fixture for creating connections

e7fd0b2

changing the simpler occurences

a51b0d0

amoghrajesh requested review from hussein-awala, eladkal and o-nikolas as code owners June 19, 2025 11:05

amoghrajesh marked this pull request as draft June 19, 2025 11:05

using env backend

d4472ec

amoghrajesh commented Jun 19, 2025

View reviewed changes

providers/asana/tests/unit/asana/operators/test_asana_tasks.py Show resolved Hide resolved

amoghrajesh requested a review from jason810496 June 20, 2025 13:31

jason810496 approved these changes Jun 20, 2025

View reviewed changes

providers/apache/kafka/tests/integration/apache/kafka/operators/test_consume.py Show resolved Hide resolved

trying to use as_json instead

27d6957

potiuk approved these changes Jun 20, 2025

View reviewed changes

amoghrajesh merged commit 8e21e7c into apache:main Jun 22, 2025
98 checks passed

amoghrajesh deleted the replace-db-merge-conn branch June 22, 2025 08:32

potiuk mentioned this pull request Jun 22, 2025

Remove pytest.mark.db_test: airbyte and amazon providers where possible #52017

Merged

potiuk mentioned this pull request Jun 22, 2025

Remove pytest.mark.db_test from providers where possible #52020

Closed

72 tasks

amoghrajesh mentioned this pull request Jun 22, 2025

Replace usage of os.environ with conf_vars in kafka integration tests #52025

Merged

dominikhei mentioned this pull request Jun 22, 2025

Added JWT authentication to the TableauHook #51756

Open

Introducing fixture to create Connections without DB in provider tests #51930

Introducing fixture to create Connections without DB in provider tests #51930

Conversation

amoghrajesh commented Jun 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why?

How?

Effort involved:

TODO

Uh oh!

potiuk commented Jun 19, 2025

Uh oh!

potiuk commented Jun 19, 2025

Uh oh!

amoghrajesh commented Jun 19, 2025

Uh oh!

potiuk commented Jun 19, 2025

Uh oh!

amoghrajesh commented Jun 19, 2025

Uh oh!

Uh oh!

amoghrajesh commented Jun 20, 2025

Uh oh!

amoghrajesh commented Jun 20, 2025

Uh oh!

jason810496 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

amoghrajesh commented Jun 20, 2025

Uh oh!

amoghrajesh commented Jun 20, 2025

Uh oh!

amoghrajesh commented Jun 20, 2025

Uh oh!

potiuk left a comment

Choose a reason for hiding this comment

Uh oh!

potiuk commented Jun 21, 2025

Uh oh!

amoghrajesh commented Jun 22, 2025

Uh oh!

Uh oh!

amoghrajesh commented Jun 22, 2025

Uh oh!

potiuk commented Jun 22, 2025

Uh oh!

potiuk commented Jun 22, 2025

Uh oh!

amoghrajesh commented Jun 23, 2025

Uh oh!

potiuk commented Jun 23, 2025

Uh oh!

amoghrajesh commented Jun 23, 2025

Uh oh!

Uh oh!

Introducing fixture to create `Connections` without DB in provider tests #51930

Introducing fixture to create `Connections` without DB in provider tests #51930

amoghrajesh commented Jun 19, 2025 •

edited

Loading