.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto/testing/mocking.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_auto_testing_mocking.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_testing_mocking.py:


Mock Tasks for Testing
--------------------------

A lot of the tasks that you write you can run locally, but some of them you will not be able to, usually because they
are tasks that depend on a third-party only available on the backend. Hive tasks are a common example, as most users
will not have access to the service that executes Hive queries from their development environment. However, it's still
useful to be able to locally run a workflow that calls such a task. In these instances, flytekit provides a couple
of utilities to help navigate this.

.. GENERATED FROM PYTHON SOURCE LINES 12-20

.. code-block:: default


    import datetime

    import pandas
    from flytekit import SQLTask, TaskMetadata, kwtypes, task, workflow
    from flytekit.testing import patch, task_mock
    from flytekit.types.schema import FlyteSchema


.. GENERATED FROM PYTHON SOURCE LINES 21-23

This is a generic SQL task (and is by default not hooked up to any datastore nor handled by any plugin), and must
be mocked.

.. GENERATED FROM PYTHON SOURCE LINES 23-32

.. code-block:: default

    sql = SQLTask(
        "my-query",
        query_template="SELECT * FROM hive.city.fact_airport_sessions WHERE ds = '{{ .Inputs.ds }}' LIMIT 10",
        inputs=kwtypes(ds=datetime.datetime),
        outputs=kwtypes(results=FlyteSchema),
        metadata=TaskMetadata(retries=2),
    )



.. GENERATED FROM PYTHON SOURCE LINES 33-34

This is a task that can run locally

.. GENERATED FROM PYTHON SOURCE LINES 34-39

.. code-block:: default

    @task
    def t1() -> datetime.datetime:
        return datetime.datetime.now()



.. GENERATED FROM PYTHON SOURCE LINES 40-41

Declare a workflow that chains these two tasks together.

.. GENERATED FROM PYTHON SOURCE LINES 41-47

.. code-block:: default

    @workflow
    def my_wf() -> FlyteSchema:
        dt = t1()
        return sql(ds=dt)



.. GENERATED FROM PYTHON SOURCE LINES 48-50

Without a mock, calling the workflow would typically raise an exception, but with the ``task_mock`` construct, which
returns a ``MagicMock`` object, we can override the return value.

.. GENERATED FROM PYTHON SOURCE LINES 50-63

.. code-block:: default

    def main_1():
        with task_mock(sql) as mock:
            mock.return_value = pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})
            assert (
                (
                    my_wf().open().all()
                    == pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})
                )
                .all()
                .all()
            )



.. GENERATED FROM PYTHON SOURCE LINES 64-66

There is another utility as well called ``patch`` which offers the same functionality, but in the traditional Python
patching style, where the first argument is the ``MagicMock`` object.

.. GENERATED FROM PYTHON SOURCE LINES 66-85

.. code-block:: default

    def main_2():
        @patch(sql)
        def test_user_demo_test(mock_sql):
            mock_sql.return_value = pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})
            assert (
                (
                    my_wf().open().all()
                    == pandas.DataFrame(data={"x": [1, 2], "y": ["3", "4"]})
                )
                .all()
                .all()
            )

        test_user_demo_test()


    if __name__ == "__main__":
        main_1()
        main_2()


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.000 seconds)


.. _sphx_glr_download_auto_testing_mocking.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example


    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: mocking.py <mocking.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: mocking.ipynb <mocking.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_