Search code examples
pythonunit-testingtestinggoogle-cloud-storagepython-mock

Unit testing Mock GCS


I have a hard time to find a way to make a unit test for the read and write methods present in this class. I am trying to create a mock with the mock patch library in order to avoid calling Google Storage but I have a hard time figure out how to do it.

from google.cloud import storage


class GCSObject(str):
    

    def __init__(self, uri=""):
        self.base, self.bucket, self.path = self.parse_uri(uri)


    def parse_uri(self, uri):
        uri = uri.lstrip("gs://").replace("//", "/").split("/", 1)
        if len(uri) > 1:
            return ("gs://", uri[0], uri[1])
        else:
            return ("gs://", uri[0], "")

    def read(self) -> bytes:
        storage_client = storage.Client()
        bucket = storage_client.bucket(self.bucket)
        return bucket.blob(self.path).download_as_string()

    def write(self, content: bytes):
        storage_client = storage.Client()
        bucket = storage_client.get_bucket(self.bucket)
        blob = bucket.blob(self.path)
        blob.upload_from_string(content)

What I am currently trying to do

@mock.patch("packages.pyGCP.pyGCP.gcs.storage.Client")
def test_upload(client):
    gcs_path = GCSObject("fake_path")
    reader = gcs_path.read()  
    #  confused what I should assert 

Solution

  • I am using google-cloud-storage==1.32.0 and python 3.7.5. Here is the unit test solution:

    gcs.py:

    from google.cloud import storage
    
    
    class GCSObject(str):
    
        def __init__(self, uri=""):
            self.base, self.bucket, self.path = self.parse_uri(uri)
    
        def parse_uri(self, uri):
            uri = uri.lstrip("gs://").replace("//", "/").split("/", 1)
            if len(uri) > 1:
                return ("gs://", uri[0], uri[1])
            else:
                return ("gs://", uri[0], "")
    
        def read(self) -> bytes:
            storage_client = storage.Client()
            bucket = storage_client.bucket(self.bucket)
            return bucket.blob(self.path).download_as_string()
    
        def write(self, content: bytes):
            storage_client = storage.Client()
            bucket = storage_client.get_bucket(self.bucket)
            blob = bucket.blob(self.path)
            blob.upload_from_string(content)
    

    test_gcs.py:

    import unittest
    from unittest import mock
    from unittest.mock import Mock
    from gcs import GCSObject
    
    
    class TestGCSObject(unittest.TestCase):
        @mock.patch('gcs.storage')
        def test_read(self, mock_storage):
            mock_gcs_client = mock_storage.Client.return_value
            mock_bucket = Mock()
            mock_bucket.blob.return_value.download_as_string.return_value = "teresa teng".encode('utf-8')
            mock_gcs_client.bucket.return_value = mock_bucket
            gcs = GCSObject('fake_path')
            actual = gcs.read()
            mock_storage.Client.assert_called_once()
            mock_gcs_client.bucket.assert_called_once_with('fake_path')
            mock_bucket.blob.assert_called_once_with('')
            mock_bucket.blob.return_value.download_as_string.assert_called_once()
            self.assertEqual(actual, "teresa teng".encode('utf-8'))
    
        @mock.patch('gcs.storage')
        def test_write(self, mock_storage):
            mock_gcs_client = mock_storage.Client.return_value
            mock_bucket = Mock()
            mock_gcs_client.get_bucket.return_value = mock_bucket
            gcs = GCSObject('fake_path')
            gcs.write(b'teresa teng')
            mock_storage.Client.assert_called_once()
            mock_gcs_client.get_bucket.assert_called_once_with('fake_path')
            mock_bucket.blob.assert_called_once_with('')
            mock_bucket.blob.return_value.upload_from_string.assert_called_once_with(b'teresa teng')
    
    
    if __name__ == '__main__':
        unittest.main()
    

    unit test result:

    ..
    ----------------------------------------------------------------------
    Ran 2 tests in 0.004s
    
    OK
    Name                                     Stmts   Miss  Cover   Missing
    ----------------------------------------------------------------------
    src/stackoverflow/64672497/gcs.py           18      1    94%   12
    src/stackoverflow/64672497/test_gcs.py      29      0   100%
    ----------------------------------------------------------------------
    TOTAL                                       47      1    98%