Python mocking table schema

I have a simple python function that needs tested that looks like this:

# Get columns from BigQuery table
def get_schema_for_table(bq_client, source, schema, table):
    bq_table = bq_client.get_table(f'{source}.{schema}.{table}')
    return bq_table.schema

However, to test this, I don’t want to test with a real table – I would like to create a Mock schema within my test (and a mock client too).

My test currently looks like this:

def test_get_schema_for_table(mocker):
    mock_client = mocker.patch('google.cloud.bigquery.Client', autospec=True)
    Table.schema = MagicMock(return_value=[
        SchemaField('column_1', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_2', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_3', 'STRING', 'NULLABLE', None, (), None)
    ])
    source = 'project'
    schema_name = 'dataset'
    table_name = 'existing_table'
    expected_schema = [SchemaField('column_1', 'STRING', 'NULLABLE', None, (), None), SchemaField('column_2', 'STRING', 'NULLABLE', None, (), None), SchemaField('column_3', 'STRING', 'NULLABLE', None, (), None)]
    returned_schema = get_schema_for_table(mock_client,source,schema_name,table_name)    
    print(returned_schema)
    assert returned_schema == expected_schema

However, when I print returned_schema, I get the following: <MagicMock name='Client.get_table().schema' id='4591155568'>.

Can anyone advise what I am doing wrong?

Answer

Well, you set the value for Table.schema but you have already mocked google.cloud.bigquery.Client. When you access any value on your mock_client you just get another magic mock – but you can define what these mocks should return.

So you need to set the expected return value on your mock_client directly.

Replace you code with this and it should work (although the test might seem a little pointless at that point)

    mock_client.get_table("project.dataset.existing_table").schema = [
        SchemaField('column_1', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_2', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_3', 'STRING', 'NULLABLE', None, (), None)
    ]

Although if you want to test more properties I’d suggest the following approach:

    mock_table = mock_client.get_table("project.dataset.existing_table")
    mock_table.schema = [
        SchemaField('column_1', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_2', 'STRING', 'NULLABLE', None, (), None),
        SchemaField('column_3', 'STRING', 'NULLABLE', None, (), None)
    ]
    # you can mock additional properties of your table