OiO.lk Community platform!

Oio.lk is an excellent forum for developers, providing a wide range of resources, discussions, and support for those in the developer community. Join oio.lk today to connect with like-minded professionals, share insights, and stay updated on the latest trends and technologies in the development field.
  You need to log in or register to access the solved answers to this problem.
  • You have reached the maximum number of guest views allowed
  • Please register below to remove this limitation

Mocking Airflow ImapHook class

  • Thread starter Thread starter Lays Rodrigues
  • Start date Start date
L

Lays Rodrigues

Guest
I wrote an Airflow Dag that uses the ImapHook to log into an email account and retrieve some mails. I'm trying to write unit tests for this dag, but I'm having issues on mocking the ImapHook class.

For the mocks to work, I need to mock the mail_client attribute that it's inside ImapHook Class, however, so far I'm unable to accomplish this mock.

The part of the code that I'm trying to mock:

Code:
    with ImapHook(imap_conn_id="imap") as hook:
        try:
            hook.mail_client.select("INBOX", readonly=True)
            _, data = hook.mail_client.search(
                "iso-8859-1",
                DATE_SEARCH_CRITERIA,
                SUBJECT_SEARCH_CRITERIA,
            )

With:

@patch("file_where_i_import_the_hook.ImapHook") I can mock the Hook, but when trying to mock the return_value of it, I get the error not enough values to unpack (expected 2, got 0) on the _, data of the search function. The mocks that I tried:

Code:
mock.return_value.select.return_value = ("OK", b"1")
mock.return_value.search.return_value = ("OK", [b"1"])

If I try this

Code:
mock.select.return_value = ("OK", b"1")
mock.search.return_value = ("OK", [b"1"])

or this

Code:
mock.mail_client.select.return_value = ("OK", b"1")
mock.mail_client.search.return_value = ("OK", [b"1"])

I get the same error. It appears that I'm unable to get to the mail_client object to mock it.

I also tried to use @patch.object to mock __enter__ and get_conn also without success. If mail_client was available on the class ImapHook, I think that the patch.objecton it would work, however if I try @patch.object(ImapHook,"mail_client") I get the error:

Code:
AttributeError: <class 'airflow.providers.imap.hooks.imap.ImapHook'> does not have the attribute 'mail_client'

The alternative would be to implement the use directly of the Imap library without using the airflow hook. However, I would like to figure out where my mock is wrong or if this is some airflow witchcraft that I need to understand.
<p>I wrote an Airflow Dag that uses the ImapHook to log into an email account and retrieve some mails. I'm trying to write unit tests for this dag, but I'm having issues on mocking the ImapHook class.</p>
<p>For the mocks to work, I need to mock the <code>mail_client</code> <a href="https://github.com/apache/airflow/blob/main/airflow/providers/imap/hooks/imap.py#L60" rel="nofollow noreferrer">attribute that it's inside ImapHook</a> Class, however, so far I'm unable to accomplish this mock.</p>
<p>The part of the code that I'm trying to mock:</p>
<pre class="lang-py prettyprint-override"><code> with ImapHook(imap_conn_id="imap") as hook:
try:
hook.mail_client.select("INBOX", readonly=True)
_, data = hook.mail_client.search(
"iso-8859-1",
DATE_SEARCH_CRITERIA,
SUBJECT_SEARCH_CRITERIA,
)
</code></pre>
<p>With:</p>
<p><code>@patch("file_where_i_import_the_hook.ImapHook")</code>
I can mock the Hook, but when trying to mock the <code>return_value</code> of it, I get the error <code>not enough values to unpack (expected 2, got 0)</code> on the <code>_, data</code> of the search function.
The mocks that I tried:</p>
<pre class="lang-py prettyprint-override"><code>mock.return_value.select.return_value = ("OK", b"1")
mock.return_value.search.return_value = ("OK", [b"1"])
</code></pre>
<p>If I try this</p>
<pre class="lang-py prettyprint-override"><code>mock.select.return_value = ("OK", b"1")
mock.search.return_value = ("OK", [b"1"])
</code></pre>
<p>or this</p>
<pre class="lang-py prettyprint-override"><code>mock.mail_client.select.return_value = ("OK", b"1")
mock.mail_client.search.return_value = ("OK", [b"1"])
</code></pre>
<p>I get the same error.
It appears that I'm unable to get to the <code>mail_client</code> object to mock it.</p>
<p>I also tried to use <code>@patch.object</code> to mock <code>__enter__</code> and <code>get_conn</code> also without success. If <code>mail_client</code> was available on the class ImapHook, I think that the <code>patch.object</code>on it would work, however if I try <code>@patch.object(ImapHook,"mail_client")</code> I get the error:</p>
<pre><code>AttributeError: <class 'airflow.providers.imap.hooks.imap.ImapHook'> does not have the attribute 'mail_client'
</code></pre>
<p>The alternative would be to implement the use directly of the Imap library without using the airflow hook. However, I would like to figure out where my mock is wrong or if this is some airflow witchcraft that I need to understand.</p>
 

Online statistics

Members online
0
Guests online
4
Total visitors
4
Top