How to mock an existing function that creates a Bigquery client

Question

I'm relatively new to pytest and unit tests in general. Specifically I'm struggling to implement mocking, but I do understand the concept. I have a Class, let's call it MyClass. It has a constructor which takes a number of arguments used by other functions within MyClass.

I have a get_tables() function that I have to test and it relies on some arguments defined in the constructor. I need to mock the BigQuery connection and return a mocked list of tables. A small snippet of the script is below

from google.cloud import bigquery
from google.cloud import storage
import logging

class MyClass:

    def __init__(self, project_id: str, pro_dataset_id: str, balancing_dataset_id: str, files: list,
                 key_file: str = None, run_local=True):
        """Creates a BigQuery client.

        Args:
            project_id:  (string), name of project
            pro_dataset_id:  (string), name of production dataset
            balancing_dataset_id:  (string), name of balancing dataset
            files:  (list), list of tables
            key_file: (string), path to the key file
        """
        if key_file is None:
            self.bq_client = bigquery.Client(project=project_id)
            self.storage_client = storage.Client(project=project_id)
            self.bucket = self.storage_client.get_bucket("{0}-my-bucket".format(project_id))

        else:
            self.bq_client = bigquery.Client.from_service_account_json(key_file)
        
        self.project_id = project_id
        self.run_local = run_local
        self.pro_dataset_id = pro_dataset_id
        self.balancing_dataset_id = balancing_dataset_id

    def get_tables(self):
        """Gets full list of all tables in a BigQuery dataset.
        Args:
        Returns:
            List of tables from a specified dataset
            """
        full_table_list = []
        dataset_ref = '{0}.{1}'.format(self.project_id, self.pro_dataset_id)
        tables = self.bq_client.list_tables(dataset_ref)  # Make an API request.
        logging.info(f"Tables contained in '{dataset_ref}':")
        for table in tables:
            full_table_list.append(table.table_id)
        logging.info(f"tables: {full_table_list}")
        return full_table_list

This was my attempt at mocking the connection and response based on an amalgamation of articles I've read and stackoverflow answers on other questions including this one How can I mock requests and the response?

import pytest
from unittest import mock
from MyPackage import MyClass

class TestMyClass:

    def mocked_list_tables():
        tables = ['mock_table1', 'mock_table2', 'mock_table3', 'mock_table4']
        return tables

    @mock.patch('MyPackage.MyClass', side_effect=mocked_list_tables)
    def test_get_tables(self):
        m = MyClass()
        assert m.get_tables() == ['mock_table1', 'mock_table2', 'mock_table3', 'mock_table4']

This is the error I get with the above unit test

TypeError: test_get_tables() takes 1 positional argument but 2 were given

How do I get this test to work correctly? Incase you're wondering, the arguments in the constructor are declared using argparse.ArgumentParser().add_argument()

I'm not in front of a computer to test this. So I'll try to explain it here best I could. You are mocking out the wrong thing. What you want to do instead is mock out the classes you are using from the `bigquery`and `storage` modules. You are doing this: `'MyPackage.MyClass'`, and I think what you should be doing instead is this: `'MyPackage.bigquery.Client'` and `MyPackage.storage.Client`. With those two you are specifically mocking out those `Client` classes so you can control the behaviour within the `test_get_tables` method you are looking to test. — idjaw, Sep 01 '21 at 23:11
To go a step further with your mocking, you should take the `spec` of the class you are mocking to strengthen your tests that much more. The spec will set the known methods in that class so your testing will ensure those methods do in fact exist when using the mock. Example: `@mock.patch('MyPackage.bigquery.Client', Mock(spec=bigquery.Client)`. — idjaw, Sep 01 '21 at 23:15
Important to note that within your `test_get_tables` you want to call the real `get_tables` method. So you need a real instance of `MyClass`. Then before you call the real `get_tables` method, you will want to do something along the lines of: `self.bq_client.list_tables.return_value = ["table1", "table2", "table9000"]`. Because, remember, if you mocked things out properly, the `self.bq_client` will now hold the mock you configured. — idjaw, Sep 01 '21 at 23:20

How to mock an existing function that creates a Bigquery client

0 Answers0