test_ml.py

# Copyright 2019 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Test cases for the firebase_admin.ml module."""

import json

import pytest

import firebase_admin
from firebase_admin import exceptions
from firebase_admin import ml
from tests import testutils


BASE_URL = 'https://firebaseml.googleapis.com/v1beta2/'
HEADER_CLIENT_KEY = 'X-FIREBASE-CLIENT'
HEADER_CLIENT_VALUE = 'fire-admin-python/{0}'.format(firebase_admin.__version__)
PROJECT_ID = 'my-project-1'

PAGE_TOKEN = 'pageToken'
NEXT_PAGE_TOKEN = 'nextPageToken'

CREATE_TIME = '2020-01-21T20:44:27.392932Z'
CREATE_TIME_MILLIS = 1579639467392

UPDATE_TIME = '2020-01-21T22:45:29.392932Z'
UPDATE_TIME_MILLIS = 1579646729392

CREATE_TIME_2 = '2020-01-21T21:44:27.392932Z'
UPDATE_TIME_2 = '2020-01-21T23:45:29.392932Z'

ETAG = '33a64df551425fcc55e4d42a148795d9f25f89d4'
MODEL_HASH = '987987a98b98798d098098e09809fc0893897'
TAG_1 = 'Tag1'
TAG_2 = 'Tag2'
TAG_3 = 'Tag3'
TAGS = [TAG_1, TAG_2]
TAGS_2 = [TAG_1, TAG_3]

MODEL_ID_1 = 'modelId1'
MODEL_NAME_1 = 'projects/{0}/models/{1}'.format(PROJECT_ID, MODEL_ID_1)
DISPLAY_NAME_1 = 'displayName1'
MODEL_JSON_1 = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_1
}
MODEL_1 = ml.Model.from_dict(MODEL_JSON_1)

MODEL_ID_2 = 'modelId2'
MODEL_NAME_2 = 'projects/{0}/models/{1}'.format(PROJECT_ID, MODEL_ID_2)
DISPLAY_NAME_2 = 'displayName2'
MODEL_JSON_2 = {
    'name': MODEL_NAME_2,
    'displayName': DISPLAY_NAME_2
}
MODEL_2 = ml.Model.from_dict(MODEL_JSON_2)

MODEL_ID_3 = 'modelId3'
MODEL_NAME_3 = 'projects/{0}/models/{1}'.format(PROJECT_ID, MODEL_ID_3)
DISPLAY_NAME_3 = 'displayName3'
MODEL_JSON_3 = {
    'name': MODEL_NAME_3,
    'displayName': DISPLAY_NAME_3
}
MODEL_3 = ml.Model.from_dict(MODEL_JSON_3)

MODEL_STATE_PUBLISHED_JSON = {
    'published': True
}
VALIDATION_ERROR_CODE = 400
VALIDATION_ERROR_MSG = 'No model format found for {0}.'.format(MODEL_ID_1)
MODEL_STATE_ERROR_JSON = {
    'validationError': {
        'code': VALIDATION_ERROR_CODE,
        'message': VALIDATION_ERROR_MSG,
    }
}

OPERATION_NAME_1 = 'projects/{0}/operations/123'.format(PROJECT_ID)
OPERATION_NOT_DONE_JSON_1 = {
    'name': OPERATION_NAME_1,
    'metadata': {
        '@type': 'type.googleapis.com/google.firebase.ml.v1beta2.ModelOperationMetadata',
        'name': 'projects/{0}/models/{1}'.format(PROJECT_ID, MODEL_ID_1),
        'basic_operation_status': 'BASIC_OPERATION_STATUS_UPLOADING'
    }
}

GCS_BUCKET_NAME = 'my_bucket'
GCS_BLOB_NAME = 'mymodel.tflite'
GCS_TFLITE_URI = 'gs://{0}/{1}'.format(GCS_BUCKET_NAME, GCS_BLOB_NAME)
GCS_TFLITE_URI_JSON = {'gcsTfliteUri': GCS_TFLITE_URI}
GCS_TFLITE_MODEL_SOURCE = ml.TFLiteGCSModelSource(GCS_TFLITE_URI)
TFLITE_FORMAT_JSON = {
    'gcsTfliteUri': GCS_TFLITE_URI,
    'sizeBytes': '1234567'
}
TFLITE_FORMAT = ml.TFLiteFormat.from_dict(TFLITE_FORMAT_JSON)

GCS_TFLITE_SIGNED_URI_PATTERN = (
    'https://storage.googleapis.com/{0}/{1}?X-Goog-Algorithm=GOOG4-RSA-SHA256&foo')
GCS_TFLITE_SIGNED_URI = GCS_TFLITE_SIGNED_URI_PATTERN.format(GCS_BUCKET_NAME, GCS_BLOB_NAME)

GCS_TFLITE_URI_2 = 'gs://my_bucket/mymodel2.tflite'
GCS_TFLITE_URI_JSON_2 = {'gcsTfliteUri': GCS_TFLITE_URI_2}
GCS_TFLITE_MODEL_SOURCE_2 = ml.TFLiteGCSModelSource(GCS_TFLITE_URI_2)
TFLITE_FORMAT_JSON_2 = {
    'gcsTfliteUri': GCS_TFLITE_URI_2,
    'sizeBytes': '2345678'
}
TFLITE_FORMAT_2 = ml.TFLiteFormat.from_dict(TFLITE_FORMAT_JSON_2)

CREATED_UPDATED_MODEL_JSON_1 = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_1,
    'createTime': CREATE_TIME,
    'updateTime': UPDATE_TIME,
    'state': MODEL_STATE_ERROR_JSON,
    'etag': ETAG,
    'modelHash': MODEL_HASH,
    'tags': TAGS,
}
CREATED_UPDATED_MODEL_1 = ml.Model.from_dict(CREATED_UPDATED_MODEL_JSON_1)

LOCKED_MODEL_JSON_1 = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_1,
    'createTime': CREATE_TIME,
    'updateTime': UPDATE_TIME,
    'tags': TAGS,
    'activeOperations': [OPERATION_NOT_DONE_JSON_1]
}

LOCKED_MODEL_JSON_2 = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_2,
    'createTime': CREATE_TIME_2,
    'updateTime': UPDATE_TIME_2,
    'tags': TAGS_2,
    'activeOperations': [OPERATION_NOT_DONE_JSON_1]
}

OPERATION_DONE_MODEL_JSON_1 = {
    'done': True,
    'response': CREATED_UPDATED_MODEL_JSON_1
}
OPERATION_MALFORMED_JSON_1 = {
    'done': True,
    # if done is true then either response or error should be populated
}
OPERATION_MISSING_NAME = {
    # Name is required if the operation is not done.
    'done': False
}
OPERATION_ERROR_CODE = 3
OPERATION_ERROR_MSG = "Invalid argument"
OPERATION_ERROR_EXPECTED_STATUS = 'INVALID_ARGUMENT'
OPERATION_ERROR_JSON_1 = {
    'done': True,
    'error': {
        'code': OPERATION_ERROR_CODE,
        'message': OPERATION_ERROR_MSG,
    }
}

FULL_MODEL_ERR_STATE_LRO_JSON = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_1,
    'createTime': CREATE_TIME,
    'updateTime': UPDATE_TIME,
    'state': MODEL_STATE_ERROR_JSON,
    'etag': ETAG,
    'modelHash': MODEL_HASH,
    'tags': TAGS,
    'activeOperations': [OPERATION_NOT_DONE_JSON_1],
}
FULL_MODEL_PUBLISHED_JSON = {
    'name': MODEL_NAME_1,
    'displayName': DISPLAY_NAME_1,
    'createTime': CREATE_TIME,
    'updateTime': UPDATE_TIME,
    'state': MODEL_STATE_PUBLISHED_JSON,
    'etag': ETAG,
    'modelHash': MODEL_HASH,
    'tags': TAGS,
    'tfliteModel': TFLITE_FORMAT_JSON
}
FULL_MODEL_PUBLISHED = ml.Model.from_dict(FULL_MODEL_PUBLISHED_JSON)
OPERATION_DONE_FULL_MODEL_PUBLISHED_JSON = {
    'name': OPERATION_NAME_1,
    'done': True,
    'response': FULL_MODEL_PUBLISHED_JSON
}

EMPTY_RESPONSE = json.dumps({})
OPERATION_NOT_DONE_RESPONSE = json.dumps(OPERATION_NOT_DONE_JSON_1)
OPERATION_DONE_RESPONSE = json.dumps(OPERATION_DONE_MODEL_JSON_1)
OPERATION_DONE_PUBLISHED_RESPONSE = json.dumps(OPERATION_DONE_FULL_MODEL_PUBLISHED_JSON)
OPERATION_ERROR_RESPONSE = json.dumps(OPERATION_ERROR_JSON_1)
OPERATION_MALFORMED_RESPONSE = json.dumps(OPERATION_MALFORMED_JSON_1)
OPERATION_MISSING_NAME_RESPONSE = json.dumps(OPERATION_MISSING_NAME)
DEFAULT_GET_RESPONSE = json.dumps(MODEL_JSON_1)
LOCKED_MODEL_2_RESPONSE = json.dumps(LOCKED_MODEL_JSON_2)
NO_MODELS_LIST_RESPONSE = json.dumps({})
DEFAULT_LIST_RESPONSE = json.dumps({
    'models': [MODEL_JSON_1, MODEL_JSON_2],
    'nextPageToken': NEXT_PAGE_TOKEN
})
LAST_PAGE_LIST_RESPONSE = json.dumps({
    'models': [MODEL_JSON_3]
})
ONE_PAGE_LIST_RESPONSE = json.dumps({
    'models': [MODEL_JSON_1, MODEL_JSON_2, MODEL_JSON_3],
})

ERROR_CODE_NOT_FOUND = 404
ERROR_MSG_NOT_FOUND = 'The resource was not found'
ERROR_STATUS_NOT_FOUND = 'NOT_FOUND'
ERROR_JSON_NOT_FOUND = {
    'error': {
        'code': ERROR_CODE_NOT_FOUND,
        'message': ERROR_MSG_NOT_FOUND,
        'status': ERROR_STATUS_NOT_FOUND
    }
}
ERROR_RESPONSE_NOT_FOUND = json.dumps(ERROR_JSON_NOT_FOUND)

ERROR_CODE_BAD_REQUEST = 400
ERROR_MSG_BAD_REQUEST = 'Invalid Argument'
ERROR_STATUS_BAD_REQUEST = 'INVALID_ARGUMENT'
ERROR_JSON_BAD_REQUEST = {
    'error': {
        'code': ERROR_CODE_BAD_REQUEST,
        'message': ERROR_MSG_BAD_REQUEST,
        'status': ERROR_STATUS_BAD_REQUEST
    }
}
ERROR_RESPONSE_BAD_REQUEST = json.dumps(ERROR_JSON_BAD_REQUEST)

INVALID_MODEL_ID_ARGS = [
    ('', ValueError),
    ('&_*#@:/?', ValueError),
    (None, TypeError),
    (12345, TypeError),
]
INVALID_MODEL_ARGS = [
    'abc',
    4.2,
    list(),
    dict(),
    True,
    -1,
    0,
    None
]
INVALID_OP_NAME_ARGS = [
    'abc',
    '123',
    'operations/project/1234/model/abc/operation/123',
    'projects/operations/123',
    'projects/$#@/operations/123',
    'projects/1234/operations/123/extrathing',
]
PAGE_SIZE_VALUE_ERROR_MSG = 'Page size must be a positive integer between ' \
                            '1 and {0}'.format(ml._MAX_PAGE_SIZE)
INVALID_STRING_OR_NONE_ARGS = [0, -1, 4.2, 0x10, False, list(), dict()]


# For validation type errors
def check_error(excinfo, err_type, msg=None):
    err = excinfo.value
    assert isinstance(err, err_type)
    if msg:
        assert str(err) == msg


# For errors that are returned in an operation
def check_operation_error(excinfo, code, msg):
    err = excinfo.value
    assert isinstance(err, exceptions.FirebaseError)
    assert err.code == code
    assert str(err) == msg


# For rpc errors
def check_firebase_error(excinfo, code, status, msg):
    err = excinfo.value
    assert isinstance(err, exceptions.FirebaseError)
    assert err.code == code
    assert err.http_response is not None
    assert err.http_response.status_code == status
    assert str(err) == msg


def instrument_ml_service(status=200, payload=None, operations=False, app=None):
    if not app:
        app = firebase_admin.get_app()
    ml_service = ml._get_ml_service(app)
    recorder = []
    session_url = 'https://firebaseml.googleapis.com/v1beta2/'

    if isinstance(status, list):
        adapter = testutils.MockMultiRequestAdapter
    else:
        adapter = testutils.MockAdapter

    if operations:
        ml_service._operation_client.session.mount(
            session_url, adapter(payload, status, recorder))
    else:
        ml_service._client.session.mount(
            session_url, adapter(payload, status, recorder))
    return recorder

class _TestStorageClient:
    @staticmethod
    def upload(bucket_name, model_file_name, app):
        del app # unused variable
        blob_name = ml._CloudStorageClient.BLOB_NAME.format(model_file_name)
        return ml._CloudStorageClient.GCS_URI.format(bucket_name, blob_name)

    @staticmethod
    def sign_uri(gcs_tflite_uri, app):
        del app # unused variable
        bucket_name, blob_name = ml._CloudStorageClient._parse_gcs_tflite_uri(gcs_tflite_uri)
        return GCS_TFLITE_SIGNED_URI_PATTERN.format(bucket_name, blob_name)

class TestModel:
    """Tests ml.Model class."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})
        ml._MLService.POLL_BASE_WAIT_TIME_SECONDS = 0.1  # shorter for test
        ml.TFLiteGCSModelSource._STORAGE_CLIENT = _TestStorageClient()

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _op_url(project_id):
        return BASE_URL + \
            'projects/{0}/operations/123'.format(project_id)

    def test_model_success_err_state_lro(self):
        model = ml.Model.from_dict(FULL_MODEL_ERR_STATE_LRO_JSON)
        assert model.model_id == MODEL_ID_1
        assert model.display_name == DISPLAY_NAME_1
        assert model.create_time == CREATE_TIME_MILLIS
        assert model.update_time == UPDATE_TIME_MILLIS
        assert model.validation_error == VALIDATION_ERROR_MSG
        assert model.published is False
        assert model.etag == ETAG
        assert model.model_hash == MODEL_HASH
        assert model.tags == TAGS
        assert model.locked is True
        assert model.model_format is None
        assert model.as_dict() == FULL_MODEL_ERR_STATE_LRO_JSON

    def test_model_success_published(self):
        model = ml.Model.from_dict(FULL_MODEL_PUBLISHED_JSON)
        assert model.model_id == MODEL_ID_1
        assert model.display_name == DISPLAY_NAME_1
        assert model.create_time == CREATE_TIME_MILLIS
        assert model.update_time == UPDATE_TIME_MILLIS
        assert model.validation_error is None
        assert model.published is True
        assert model.etag == ETAG
        assert model.model_hash == MODEL_HASH
        assert model.tags == TAGS
        assert model.locked is False
        assert model.model_format == TFLITE_FORMAT
        assert model.as_dict() == FULL_MODEL_PUBLISHED_JSON

    def test_model_keyword_based_creation_and_setters(self):
        model = ml.Model(display_name=DISPLAY_NAME_1, tags=TAGS, model_format=TFLITE_FORMAT)
        assert model.display_name == DISPLAY_NAME_1
        assert model.tags == TAGS
        assert model.model_format == TFLITE_FORMAT
        assert model.as_dict() == {
            'displayName': DISPLAY_NAME_1,
            'tags': TAGS,
            'tfliteModel': TFLITE_FORMAT_JSON
        }

        model.display_name = DISPLAY_NAME_2
        model.tags = TAGS_2
        model.model_format = TFLITE_FORMAT_2
        assert model.as_dict() == {
            'displayName': DISPLAY_NAME_2,
            'tags': TAGS_2,
            'tfliteModel': TFLITE_FORMAT_JSON_2
        }

    def test_gcs_tflite_model_format_source_creation(self):
        model_source = ml.TFLiteGCSModelSource(gcs_tflite_uri=GCS_TFLITE_URI)
        model_format = ml.TFLiteFormat(model_source=model_source)
        model = ml.Model(display_name=DISPLAY_NAME_1, model_format=model_format)
        assert model.as_dict() == {
            'displayName': DISPLAY_NAME_1,
            'tfliteModel': {
                'gcsTfliteUri': GCS_TFLITE_URI
            }
        }

    def test_source_creation_from_tflite_file(self):
        model_source = ml.TFLiteGCSModelSource.from_tflite_model_file(
            "my_model.tflite", "my_bucket")
        assert model_source.as_dict() == {
            'gcsTfliteUri': 'gs://my_bucket/Firebase/ML/Models/my_model.tflite'
        }

    def test_gcs_tflite_model_source_setters(self):
        model_source = ml.TFLiteGCSModelSource(GCS_TFLITE_URI)
        model_source.gcs_tflite_uri = GCS_TFLITE_URI_2
        assert model_source.gcs_tflite_uri == GCS_TFLITE_URI_2
        assert model_source.as_dict() == GCS_TFLITE_URI_JSON_2

    def test_model_format_setters(self):
        model_format = ml.TFLiteFormat(model_source=GCS_TFLITE_MODEL_SOURCE)
        model_format.model_source = GCS_TFLITE_MODEL_SOURCE_2
        assert model_format.model_source == GCS_TFLITE_MODEL_SOURCE_2
        assert model_format.as_dict() == {
            'tfliteModel': {
                'gcsTfliteUri': GCS_TFLITE_URI_2
            }
        }

    def test_model_as_dict_for_upload(self):
        model_source = ml.TFLiteGCSModelSource(gcs_tflite_uri=GCS_TFLITE_URI)
        model_format = ml.TFLiteFormat(model_source=model_source)
        model = ml.Model(display_name=DISPLAY_NAME_1, model_format=model_format)
        assert model.as_dict(for_upload=True) == {
            'displayName': DISPLAY_NAME_1,
            'tfliteModel': {
                'gcsTfliteUri': GCS_TFLITE_SIGNED_URI
            }
        }

    @pytest.mark.parametrize('helper_func', [
        ml.TFLiteGCSModelSource.from_keras_model,
        ml.TFLiteGCSModelSource.from_saved_model
    ])
    def test_tf_not_enabled(self, helper_func):
        ml._TF_ENABLED = False # for reliability
        with pytest.raises(ImportError) as excinfo:
            helper_func(None)
        check_error(excinfo, ImportError)

    @pytest.mark.parametrize('display_name, exc_type', [
        ('', ValueError),
        ('&_*#@:/?', ValueError),
        (12345, TypeError)
    ])
    def test_model_display_name_validation_errors(self, display_name, exc_type):
        with pytest.raises(exc_type) as excinfo:
            ml.Model(display_name=display_name)
        check_error(excinfo, exc_type)

    @pytest.mark.parametrize('tags, exc_type, error_message', [
        ('tag1', TypeError, 'Tags must be a list of strings.'),
        (123, TypeError, 'Tags must be a list of strings.'),
        (['tag1', 123, 'tag2'], TypeError, 'Tags must be a list of strings.'),
        (['tag1', '@#$%^&'], ValueError, 'Tag format is invalid.'),
        (['', 'tag2'], ValueError, 'Tag format is invalid.'),
        (['sixty-one_characters_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx',
          'tag2'], ValueError, 'Tag format is invalid.')
    ])
    def test_model_tags_validation_errors(self, tags, exc_type, error_message):
        with pytest.raises(exc_type) as excinfo:
            ml.Model(tags=tags)
        check_error(excinfo, exc_type, error_message)

    @pytest.mark.parametrize('model_format', [
        123,
        "abc",
        {},
        [],
        True
    ])
    def test_model_format_validation_errors(self, model_format):
        with pytest.raises(TypeError) as excinfo:
            ml.Model(model_format=model_format)
        check_error(excinfo, TypeError, 'Model format must be a ModelFormat object.')

    @pytest.mark.parametrize('model_source', [
        123,
        "abc",
        {},
        [],
        True
    ])
    def test_model_source_validation_errors(self, model_source):
        with pytest.raises(TypeError) as excinfo:
            ml.TFLiteFormat(model_source=model_source)
        check_error(excinfo, TypeError, 'Model source must be a TFLiteModelSource object.')

    @pytest.mark.parametrize('uri, exc_type', [
        (123, TypeError),
        ('abc', ValueError),
        ('gs://NO_CAPITALS', ValueError),
        ('gs://abc/', ValueError),
        ('gs://aa/model.tflite', ValueError),
        ('gs://@#$%/model.tflite', ValueError),
        ('gs://invalid space/model.tflite', ValueError),
        ('gs://sixty-four-characters_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx/model.tflite',
         ValueError)
    ])
    def test_gcs_tflite_source_validation_errors(self, uri, exc_type):
        with pytest.raises(exc_type) as excinfo:
            ml.TFLiteGCSModelSource(gcs_tflite_uri=uri)
        check_error(excinfo, exc_type)

    def test_wait_for_unlocked_not_locked(self):
        model = ml.Model(display_name="not_locked")
        model.wait_for_unlocked()

    def test_wait_for_unlocked(self):
        recorder = instrument_ml_service(status=200,
                                         operations=True,
                                         payload=OPERATION_DONE_PUBLISHED_RESPONSE)
        model = ml.Model.from_dict(LOCKED_MODEL_JSON_1)
        model.wait_for_unlocked()
        assert model == FULL_MODEL_PUBLISHED
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == TestModel._op_url(PROJECT_ID)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_wait_for_unlocked_timeout(self):
        recorder = instrument_ml_service(
            status=200, operations=True, payload=OPERATION_NOT_DONE_RESPONSE)
        ml._MLService.POLL_BASE_WAIT_TIME_SECONDS = 3 # longer so timeout applies immediately
        model = ml.Model.from_dict(LOCKED_MODEL_JSON_1)
        with pytest.raises(Exception) as excinfo:
            model.wait_for_unlocked(max_time_seconds=0.1)
        check_error(excinfo, exceptions.DeadlineExceededError, 'Polling max time exceeded.')
        assert len(recorder) == 1


class TestCreateModel:
    """Tests ml.create_model."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})
        ml._MLService.POLL_BASE_WAIT_TIME_SECONDS = 0.1  # shorter for test

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _url(project_id):
        return BASE_URL + 'projects/{0}/models'.format(project_id)

    @staticmethod
    def _op_url(project_id):
        return BASE_URL + \
            'projects/{0}/operations/123'.format(project_id)

    @staticmethod
    def _get_url(project_id, model_id):
        return BASE_URL + 'projects/{0}/models/{1}'.format(project_id, model_id)

    def test_immediate_done(self):
        instrument_ml_service(status=200, payload=OPERATION_DONE_RESPONSE)
        model = ml.create_model(MODEL_1)
        assert model == CREATED_UPDATED_MODEL_1

    def test_returns_locked(self):
        recorder = instrument_ml_service(
            status=[200, 200],
            payload=[OPERATION_NOT_DONE_RESPONSE, LOCKED_MODEL_2_RESPONSE])
        expected_model = ml.Model.from_dict(LOCKED_MODEL_JSON_2)
        model = ml.create_model(MODEL_1)

        assert model == expected_model
        assert len(recorder) == 2
        assert recorder[0].method == 'POST'
        assert recorder[0].url == TestCreateModel._url(PROJECT_ID)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        assert recorder[1].method == 'GET'
        assert recorder[1].url == TestCreateModel._get_url(PROJECT_ID, MODEL_ID_1)
        assert recorder[1].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_operation_error(self):
        instrument_ml_service(status=200, payload=OPERATION_ERROR_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.create_model(MODEL_1)
        # The http request succeeded, the operation returned contains a create failure
        check_operation_error(excinfo, OPERATION_ERROR_EXPECTED_STATUS, OPERATION_ERROR_MSG)

    def test_malformed_operation(self):
        instrument_ml_service(status=200, payload=OPERATION_MALFORMED_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.create_model(MODEL_1)
        check_error(excinfo, exceptions.UnknownError, 'Internal Error: Malformed Operation.')

    def test_rpc_error_create(self):
        create_recorder = instrument_ml_service(
            status=400, payload=ERROR_RESPONSE_BAD_REQUEST)
        with pytest.raises(Exception) as excinfo:
            ml.create_model(MODEL_1)
        check_firebase_error(
            excinfo,
            ERROR_STATUS_BAD_REQUEST,
            ERROR_CODE_BAD_REQUEST,
            ERROR_MSG_BAD_REQUEST
        )
        assert len(create_recorder) == 1

    @pytest.mark.parametrize('model', INVALID_MODEL_ARGS)
    def test_not_model(self, model):
        with pytest.raises(Exception) as excinfo:
            ml.create_model(model)
        check_error(excinfo, TypeError, 'Model must be an ml.Model.')

    def test_missing_display_name(self):
        with pytest.raises(Exception) as excinfo:
            ml.create_model(ml.Model.from_dict({}))
        check_error(excinfo, ValueError, 'Model must have a display name.')

    def test_missing_op_name(self):
        instrument_ml_service(status=200, payload=OPERATION_MISSING_NAME_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.create_model(MODEL_1)
        check_error(excinfo, TypeError)

    @pytest.mark.parametrize('op_name', INVALID_OP_NAME_ARGS)
    def test_invalid_op_name(self, op_name):
        payload = json.dumps({'name': op_name})
        instrument_ml_service(status=200, payload=payload)
        with pytest.raises(Exception) as excinfo:
            ml.create_model(MODEL_1)
        check_error(excinfo, ValueError, 'Operation name format is invalid.')


class TestUpdateModel:
    """Tests ml.update_model."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})
        ml._MLService.POLL_BASE_WAIT_TIME_SECONDS = 0.1  # shorter for test

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _url(project_id, model_id):
        return BASE_URL + 'projects/{0}/models/{1}'.format(project_id, model_id)

    @staticmethod
    def _op_url(project_id):
        return BASE_URL + \
            'projects/{0}/operations/123'.format(project_id)

    def test_immediate_done(self):
        instrument_ml_service(status=200, payload=OPERATION_DONE_RESPONSE)
        model = ml.update_model(MODEL_1)
        assert model == CREATED_UPDATED_MODEL_1

    def test_returns_locked(self):
        recorder = instrument_ml_service(
            status=[200, 200],
            payload=[OPERATION_NOT_DONE_RESPONSE, LOCKED_MODEL_2_RESPONSE])
        expected_model = ml.Model.from_dict(LOCKED_MODEL_JSON_2)
        model = ml.update_model(MODEL_1)

        assert model == expected_model
        assert len(recorder) == 2
        assert recorder[0].method == 'PATCH'
        assert recorder[0].url == TestUpdateModel._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        assert recorder[1].method == 'GET'
        assert recorder[1].url == TestUpdateModel._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[1].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_operation_error(self):
        instrument_ml_service(status=200, payload=OPERATION_ERROR_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.update_model(MODEL_1)
        # The http request succeeded, the operation returned contains an update failure
        check_operation_error(excinfo, OPERATION_ERROR_EXPECTED_STATUS, OPERATION_ERROR_MSG)

    def test_malformed_operation(self):
        instrument_ml_service(status=200, payload=OPERATION_MALFORMED_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.update_model(MODEL_1)
        check_error(excinfo, exceptions.UnknownError, 'Internal Error: Malformed Operation.')

    def test_rpc_error(self):
        create_recorder = instrument_ml_service(
            status=400, payload=ERROR_RESPONSE_BAD_REQUEST)
        with pytest.raises(Exception) as excinfo:
            ml.update_model(MODEL_1)
        check_firebase_error(
            excinfo,
            ERROR_STATUS_BAD_REQUEST,
            ERROR_CODE_BAD_REQUEST,
            ERROR_MSG_BAD_REQUEST
        )
        assert len(create_recorder) == 1

    @pytest.mark.parametrize('model', INVALID_MODEL_ARGS)
    def test_not_model(self, model):
        with pytest.raises(Exception) as excinfo:
            ml.update_model(model)
        check_error(excinfo, TypeError, 'Model must be an ml.Model.')

    def test_missing_display_name(self):
        with pytest.raises(Exception) as excinfo:
            ml.update_model(ml.Model.from_dict({}))
        check_error(excinfo, ValueError, 'Model must have a display name.')

    def test_missing_op_name(self):
        instrument_ml_service(status=200, payload=OPERATION_MISSING_NAME_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            ml.update_model(MODEL_1)
        check_error(excinfo, TypeError)

    @pytest.mark.parametrize('op_name', INVALID_OP_NAME_ARGS)
    def test_invalid_op_name(self, op_name):
        payload = json.dumps({'name': op_name})
        instrument_ml_service(status=200, payload=payload)
        with pytest.raises(Exception) as excinfo:
            ml.update_model(MODEL_1)
        check_error(excinfo, ValueError, 'Operation name format is invalid.')


class TestPublishUnpublish:
    """Tests ml.publish_model and ml.unpublish_model."""

    PUBLISH_UNPUBLISH_WITH_ARGS = [
        (ml.publish_model, True),
        (ml.unpublish_model, False)
    ]
    PUBLISH_UNPUBLISH_FUNCS = [item[0] for item in PUBLISH_UNPUBLISH_WITH_ARGS]

    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})
        ml._MLService.POLL_BASE_WAIT_TIME_SECONDS = 0.1  # shorter for test

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _update_url(project_id, model_id):
        update_url = 'projects/{0}/models/{1}?updateMask=state.published'.format(
            project_id, model_id)
        return BASE_URL + update_url

    @staticmethod
    def _get_url(project_id, model_id):
        return BASE_URL + 'projects/{0}/models/{1}'.format(project_id, model_id)

    @staticmethod
    def _op_url(project_id):
        return BASE_URL + \
            'projects/{0}/operations/123'.format(project_id)

    @pytest.mark.parametrize('publish_function, published', PUBLISH_UNPUBLISH_WITH_ARGS)
    def test_immediate_done(self, publish_function, published):
        recorder = instrument_ml_service(status=200, payload=OPERATION_DONE_RESPONSE)
        model = publish_function(MODEL_ID_1)
        assert model == CREATED_UPDATED_MODEL_1
        assert len(recorder) == 1
        assert recorder[0].method == 'PATCH'
        assert recorder[0].url == TestPublishUnpublish._update_url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        body = json.loads(recorder[0].body.decode())
        assert body.get('state', {}).get('published', None) is published

    @pytest.mark.parametrize('publish_function', PUBLISH_UNPUBLISH_FUNCS)
    def test_returns_locked(self, publish_function):
        recorder = instrument_ml_service(
            status=[200, 200],
            payload=[OPERATION_NOT_DONE_RESPONSE, LOCKED_MODEL_2_RESPONSE])
        expected_model = ml.Model.from_dict(LOCKED_MODEL_JSON_2)
        model = publish_function(MODEL_ID_1)

        assert model == expected_model
        assert len(recorder) == 2
        assert recorder[0].method == 'PATCH'
        assert recorder[0].url == TestPublishUnpublish._update_url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        assert recorder[1].method == 'GET'
        assert recorder[1].url == TestPublishUnpublish._get_url(PROJECT_ID, MODEL_ID_1)
        assert recorder[1].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    @pytest.mark.parametrize('publish_function', PUBLISH_UNPUBLISH_FUNCS)
    def test_operation_error(self, publish_function):
        instrument_ml_service(status=200, payload=OPERATION_ERROR_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            publish_function(MODEL_ID_1)
        # The http request succeeded, the operation returned contains an update failure
        check_operation_error(excinfo, OPERATION_ERROR_EXPECTED_STATUS, OPERATION_ERROR_MSG)

    @pytest.mark.parametrize('publish_function', PUBLISH_UNPUBLISH_FUNCS)
    def test_malformed_operation(self, publish_function):
        instrument_ml_service(status=200, payload=OPERATION_MALFORMED_RESPONSE)
        with pytest.raises(Exception) as excinfo:
            publish_function(MODEL_ID_1)
        check_error(excinfo, exceptions.UnknownError, 'Internal Error: Malformed Operation.')

    @pytest.mark.parametrize('publish_function', PUBLISH_UNPUBLISH_FUNCS)
    def test_rpc_error(self, publish_function):
        create_recorder = instrument_ml_service(
            status=400, payload=ERROR_RESPONSE_BAD_REQUEST)
        with pytest.raises(Exception) as excinfo:
            publish_function(MODEL_ID_1)
        check_firebase_error(
            excinfo,
            ERROR_STATUS_BAD_REQUEST,
            ERROR_CODE_BAD_REQUEST,
            ERROR_MSG_BAD_REQUEST
        )
        assert len(create_recorder) == 1


class TestGetModel:
    """Tests ml.get_model."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _url(project_id, model_id):
        return BASE_URL + 'projects/{0}/models/{1}'.format(project_id, model_id)

    def test_get_model(self):
        recorder = instrument_ml_service(status=200, payload=DEFAULT_GET_RESPONSE)
        model = ml.get_model(MODEL_ID_1)
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == TestGetModel._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        assert model == MODEL_1
        assert model.model_id == MODEL_ID_1
        assert model.display_name == DISPLAY_NAME_1

    @pytest.mark.parametrize('model_id, exc_type', INVALID_MODEL_ID_ARGS)
    def test_get_model_validation_errors(self, model_id, exc_type):
        with pytest.raises(exc_type) as excinfo:
            ml.get_model(model_id)
        check_error(excinfo, exc_type)

    def test_get_model_error(self):
        recorder = instrument_ml_service(status=404, payload=ERROR_RESPONSE_NOT_FOUND)
        with pytest.raises(exceptions.NotFoundError) as excinfo:
            ml.get_model(MODEL_ID_1)
        check_firebase_error(
            excinfo,
            ERROR_STATUS_NOT_FOUND,
            ERROR_CODE_NOT_FOUND,
            ERROR_MSG_NOT_FOUND
        )
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == TestGetModel._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_no_project_id(self):
        def evaluate():
            app = firebase_admin.initialize_app(testutils.MockCredential(), name='no_project_id')
            with pytest.raises(ValueError):
                ml.get_model(MODEL_ID_1, app)
        testutils.run_without_project_id(evaluate)


class TestDeleteModel:
    """Tests ml.delete_model."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _url(project_id, model_id):
        return BASE_URL + 'projects/{0}/models/{1}'.format(project_id, model_id)

    def test_delete_model(self):
        recorder = instrument_ml_service(status=200, payload=EMPTY_RESPONSE)
        ml.delete_model(MODEL_ID_1) # no response for delete
        assert len(recorder) == 1
        assert recorder[0].method == 'DELETE'
        assert recorder[0].url == TestDeleteModel._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    @pytest.mark.parametrize('model_id, exc_type', INVALID_MODEL_ID_ARGS)
    def test_delete_model_validation_errors(self, model_id, exc_type):
        with pytest.raises(exc_type) as excinfo:
            ml.delete_model(model_id)
        check_error(excinfo, exc_type)

    def test_delete_model_error(self):
        recorder = instrument_ml_service(status=404, payload=ERROR_RESPONSE_NOT_FOUND)
        with pytest.raises(exceptions.NotFoundError) as excinfo:
            ml.delete_model(MODEL_ID_1)
        check_firebase_error(
            excinfo,
            ERROR_STATUS_NOT_FOUND,
            ERROR_CODE_NOT_FOUND,
            ERROR_MSG_NOT_FOUND
        )
        assert len(recorder) == 1
        assert recorder[0].method == 'DELETE'
        assert recorder[0].url == self._url(PROJECT_ID, MODEL_ID_1)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_no_project_id(self):
        def evaluate():
            app = firebase_admin.initialize_app(testutils.MockCredential(), name='no_project_id')
            with pytest.raises(ValueError):
                ml.delete_model(MODEL_ID_1, app)
        testutils.run_without_project_id(evaluate)


class TestListModels:
    """Tests ml.list_models."""
    @classmethod
    def setup_class(cls):
        cred = testutils.MockCredential()
        firebase_admin.initialize_app(cred, {'projectId': PROJECT_ID})

    @classmethod
    def teardown_class(cls):
        testutils.cleanup_apps()

    @staticmethod
    def _url(project_id):
        return BASE_URL + 'projects/{0}/models'.format(project_id)

    @staticmethod
    def _check_page(page, model_count):
        assert isinstance(page, ml.ListModelsPage)
        assert len(page.models) == model_count
        for model in page.models:
            assert isinstance(model, ml.Model)

    def test_list_models_no_args(self):
        recorder = instrument_ml_service(status=200, payload=DEFAULT_LIST_RESPONSE)
        models_page = ml.list_models()
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == TestListModels._url(PROJECT_ID)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        TestListModels._check_page(models_page, 2)
        assert models_page.has_next_page
        assert models_page.next_page_token == NEXT_PAGE_TOKEN
        assert models_page.models[0] == MODEL_1
        assert models_page.models[1] == MODEL_2

    def test_list_models_with_all_args(self):
        recorder = instrument_ml_service(status=200, payload=LAST_PAGE_LIST_RESPONSE)
        models_page = ml.list_models(
            'display_name=displayName3',
            page_size=10,
            page_token=PAGE_TOKEN)
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == (
            TestListModels._url(PROJECT_ID) +
            '?filter=display_name%3DdisplayName3&page_size=10&page_token={0}'
            .format(PAGE_TOKEN))
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE
        assert isinstance(models_page, ml.ListModelsPage)
        assert len(models_page.models) == 1
        assert models_page.models[0] == MODEL_3
        assert not models_page.has_next_page

    @pytest.mark.parametrize('list_filter', INVALID_STRING_OR_NONE_ARGS)
    def test_list_models_list_filter_validation(self, list_filter):
        with pytest.raises(TypeError) as excinfo:
            ml.list_models(list_filter=list_filter)
        check_error(excinfo, TypeError, 'List filter must be a string or None.')

    @pytest.mark.parametrize('page_size, exc_type, error_message', [
        ('abc', TypeError, 'Page size must be a number or None.'),
        (4.2, TypeError, 'Page size must be a number or None.'),
        (list(), TypeError, 'Page size must be a number or None.'),
        (dict(), TypeError, 'Page size must be a number or None.'),
        (True, TypeError, 'Page size must be a number or None.'),
        (-1, ValueError, PAGE_SIZE_VALUE_ERROR_MSG),
        (0, ValueError, PAGE_SIZE_VALUE_ERROR_MSG),
        (ml._MAX_PAGE_SIZE + 1, ValueError, PAGE_SIZE_VALUE_ERROR_MSG)
    ])
    def test_list_models_page_size_validation(self, page_size, exc_type, error_message):
        with pytest.raises(exc_type) as excinfo:
            ml.list_models(page_size=page_size)
        check_error(excinfo, exc_type, error_message)

    @pytest.mark.parametrize('page_token', INVALID_STRING_OR_NONE_ARGS)
    def test_list_models_page_token_validation(self, page_token):
        with pytest.raises(TypeError) as excinfo:
            ml.list_models(page_token=page_token)
        check_error(excinfo, TypeError, 'Page token must be a string or None.')

    def test_list_models_error(self):
        recorder = instrument_ml_service(status=400, payload=ERROR_RESPONSE_BAD_REQUEST)
        with pytest.raises(exceptions.InvalidArgumentError) as excinfo:
            ml.list_models()
        check_firebase_error(
            excinfo,
            ERROR_STATUS_BAD_REQUEST,
            ERROR_CODE_BAD_REQUEST,
            ERROR_MSG_BAD_REQUEST
        )
        assert len(recorder) == 1
        assert recorder[0].method == 'GET'
        assert recorder[0].url == TestListModels._url(PROJECT_ID)
        assert recorder[0].headers[HEADER_CLIENT_KEY] == HEADER_CLIENT_VALUE

    def test_no_project_id(self):
        def evaluate():
            app = firebase_admin.initialize_app(testutils.MockCredential(), name='no_project_id')
            with pytest.raises(ValueError):
                ml.list_models(app=app)
        testutils.run_without_project_id(evaluate)

    def test_list_single_page(self):
        recorder = instrument_ml_service(status=200, payload=LAST_PAGE_LIST_RESPONSE)
        models_page = ml.list_models()
        assert len(recorder) == 1
        assert models_page.next_page_token == ''
        assert models_page.has_next_page is False
        assert models_page.get_next_page() is None
        models = [model for model in models_page.iterate_all()]
        assert len(models) == 1

    def test_list_multiple_pages(self):
        # Page 1
        recorder = instrument_ml_service(status=200, payload=DEFAULT_LIST_RESPONSE)
        page = ml.list_models()
        assert len(recorder) == 1
        assert len(page.models) == 2
        assert page.next_page_token == NEXT_PAGE_TOKEN
        assert page.has_next_page is True

        # Page 2
        recorder = instrument_ml_service(status=200, payload=LAST_PAGE_LIST_RESPONSE)
        page_2 = page.get_next_page()
        assert len(recorder) == 1
        assert len(page_2.models) == 1
        assert page_2.next_page_token == ''
        assert page_2.has_next_page is False
        assert page_2.get_next_page() is None

    def test_list_models_paged_iteration(self):
        # Page 1
        recorder = instrument_ml_service(status=200, payload=DEFAULT_LIST_RESPONSE)
        page = ml.list_models()
        assert page.next_page_token == NEXT_PAGE_TOKEN
        assert page.has_next_page is True
        iterator = page.iterate_all()
        for index in range(2):
            model = next(iterator)
            assert model.display_name == 'displayName{0}'.format(index+1)
        assert len(recorder) == 1

        # Page 2
        recorder = instrument_ml_service(status=200, payload=LAST_PAGE_LIST_RESPONSE)
        model = next(iterator)
        assert model.display_name == DISPLAY_NAME_3
        with pytest.raises(StopIteration):
            next(iterator)

    def test_list_models_stop_iteration(self):
        recorder = instrument_ml_service(status=200, payload=ONE_PAGE_LIST_RESPONSE)
        page = ml.list_models()
        assert len(recorder) == 1
        assert len(page.models) == 3
        iterator = page.iterate_all()
        models = [model for model in iterator]
        assert len(page.models) == 3
        with pytest.raises(StopIteration):
            next(iterator)
        assert len(models) == 3

    def test_list_models_no_models(self):
        recorder = instrument_ml_service(status=200, payload=NO_MODELS_LIST_RESPONSE)
        page = ml.list_models()
        assert len(recorder) == 1
        assert len(page.models) == 0
        models = [model for model in page.iterate_all()]
        assert len(models) == 0