Search code examples
pythongoogle-cloud-platformpython-unittest

Python Mocking / Patching multiple nested functions / variables


I am new to Python and GCP but I am trying to create some tests for my GCF function that moves a file from one bucket to another.

Simplified Python code:

import functions_framework
from google.cloud import storage

storageClient = storage.Client()

@functions_framework.cloud_event
def storage_trigger(cloud_event):
    data = cloud_event.data
    return move_file_to_folder(data)

def move_file_to_folder(data):
    file_name = data['name']
    folder_name = 'test-folder/'
    source_bucket = storageClient.get_bucket(data['bucket'])
    source_blob = source_bucket.get_blob(file_name)
    dest_file_name = file_name

    try:
        if source_bucket.get_blob(folder_name + file_name):
            #This is what I need to mock
            dest_file_name = increment_file_name(source_bucket, folder_name, file_name)

        source_bucket.copy_blob(source_blob, source_bucket, folder_name + dest_file_name)

        return f'{dest_file_name} successfully sent and moved!'
    except Exception as e:
        print(e)
        return Exception


def increment_file_name(source_bucket, folder_name, file_name):
    new_name = file_name   
    split_name = file_name.split('.')
    first_part = '.'.join(split_name[:-1])
    ext = split_name[-1]

    num = 1
    while source_bucket.get_blob(folder_name + new_name):
        new_name = f'{first_part}-{num}.{ext}'
        print(new_name)
        num += 1

    return new_name

Test

import unittest
from unittest.mock import patch, Mock
from cloudevents.http import CloudEvent
import main


attributes = {
    "id": "5e9f24a",
    "type": "google.cloud.storage.object.v1.finalized",
    "source": "sourceUrlHere",
}

data = {
        'bucket': 'test_bucket',
        'contentType': 'audio/mpeg',
        'id': 'test.mp3',
        'kind': 'storage#object',
        'name': 'test.mp3',
        'size': '184162',
        'storageClass': 'STANDARD',
        'timeCreated': '2023-07-25T13:42:25.541Z',
        'updated': '2023-07-25T13:42:25.541Z'
    }

event = CloudEvent(attributes, data)

class SuccessTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        return super().setUpClass()
    
    @patch('main.storageClient')
    @patch('main.increment_file_name')
    def test_storage_trigger(self, mock_storage, func):
        mock_storage_client = mock_storage.Client.return_value
        mock_bucket = Mock()
        mock_storage_client.get_bucket.return_value = mock_bucket

        func.return_value = 'new-test-file-name.mp3'

        expected = 'Data successfully inserted into BigQuery table.', 200
        print(main.storage_trigger(event))
        self.assertEqual(main.storage_trigger(event), expected)


if __name__ == "__main__":
    unittest.main()

Issues

If I don't mock the increment_file_name function, it just increments until eternity (I assume because the the mock is designed to just say it exists?). So I am trying to patch that function

Several resources indicate that I need to patch where the function is called not the function itself, but the function call is being assigned to a variable. I tried: @patch('main.move_file_to_folder.increment_file_name') but then it cant find the attribute increment_file_name

Question

So how do I tell the test that make the reassignment of dest_file_name to be 'new-test-file-name.mp3'?


Solution

  • I believe it will work if you use it with __name__?

    @patch('main.storageClient')
    def test_storage_trigger(self, mock_storage):
        mock_storage_client = mock_storage.Client.return_value
        mock_bucket = Mock()
        mock_storage_client.get_bucket.return_value = mock_bucket
        
    
        expected = 'Data successfully inserted into BigQuery table.', 200
        with mock.patch(
            f'{__name__}.main.increment_file_name', 
            return_value='new-test-file-name.mp3'
        ):
            print(main.storage_trigger(event))
            self.assertEqual(main.storage_trigger(event), expected)