Unittest a function that wraps a custom Airflow plugin

I implemented a custom Airflow plugin (CustomHook). Airflow has no issue importing and using my CustomHook plugin.

However, I am having issues writing a unit test for the function that wraps the plugin. The issue is an ImportError for CustomHook. This error makes sense to me because Airflow has some weird interface that does plugin "injection" (if that's the right word). The unittest test package (pytest) does not know about the custom plugin being injected by Airflow.

Any thoughts on how to test the function my_on_failure_callback?

Here is the directory structure.

airflow_home
├── dags
│   └── utils
│       └── callbacks.py
├── plugins
│   └── plugins.py
│   └── custom_hooks.py   
├── tests
│   └── utils
│       └── test_callbacks.py

Here is the function. I included the interface and CustomHook implementations at the end for reference.

# ./dags/utils/callbacks.py

from airflow.hooks.custom import CustomHook


def my_on_failure_callback(context):
    dag = context['dag']

    message = f'''
        *DAG Failure*
        ----
        *dag*: {dag.dag_id}
    '''

    hook = CustomHook(
        http_conn_id='my_connection_id',
        message=message,
    )
    hook.execute()

Here's my unittest that results in an ImportError:

import mock

from dags.utils.callbacks import my_on_failure_callback


@mock.patch('dags.utils.callbacks.CustomHook')
def test_google_chat_failure_alert(mock_hook):
   ...etc...

Results in the import error below.

ImportError while importing test module '/Users/myusser/airflow_home/tests/utils/test_callbacks.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
tests/utils/test_callbacks.py:3: in <module>
    from dags.utils.callbacks import my_on_failure_callback
dags/utils/callbacks.py:3: in <module>
    from airflow.hooks.custom import CustomHook
E   ModuleNotFoundError: No module named 'airflow.hooks.custom'

CustomHook and interface implementations for reference.

# ./plugins/custom_hooks.py

import json

from airflow.hooks.http_hook import HttpHook


class CustomHook(HttpHook):
    def __init__(self, http_conn_id, message, *args, **kwargs):
        super(CustomHook, self).__init__(http_conn_id=http_conn_id, *args, **kwargs)
        conn = self.get_connection(http_conn_id)
        self.webhook_url = conn.host
        self.message = message

    def execute(self):
        self.run(
                endpoint=self.webhook_url,
                data=json.dumps({'text': self.message}),
                headers={'Content-Type': 'application/json'}
        )

# ./plugins/plugins.py

from airflow.plugins_manager import AirflowPlugin

from custom_hooks import CustomHook


class CustomPlugins(AirflowPlugin):
    name = 'custom'

    hooks = [
       CustomHook,
    ]