-
Notifications
You must be signed in to change notification settings - Fork 339
Integration tests for Firebase ML #394
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,373 @@ | ||
# Copyright 2020 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. | ||
|
||
"""Integration tests for firebase_admin.ml module.""" | ||
import os | ||
import random | ||
import re | ||
import shutil | ||
import string | ||
import tempfile | ||
import pytest | ||
|
||
|
||
from firebase_admin import exceptions | ||
from firebase_admin import ml | ||
from tests import testutils | ||
|
||
|
||
# pylint: disable=import-error,no-name-in-module | ||
try: | ||
import tensorflow as tf | ||
_TF_ENABLED = True | ||
except ImportError: | ||
_TF_ENABLED = False | ||
|
||
|
||
def _random_identifier(prefix): | ||
#pylint: disable=unused-variable | ||
suffix = ''.join([random.choice(string.ascii_letters + string.digits) for n in range(8)]) | ||
return '{0}_{1}'.format(prefix, suffix) | ||
|
||
|
||
NAME_ONLY_ARGS = { | ||
'display_name': _random_identifier('TestModel123_') | ||
} | ||
NAME_ONLY_ARGS_UPDATED = { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Now that we have done the refactor, I wonder if it makes sense to define each of these as a fixture.
That simplifies the call-site significantly, and makes sure multiple calls to the same fixture don't reuse identifiers. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We started with something very similar to this. I think we've done enough refactoring now. |
||
'display_name': _random_identifier('TestModel123_updated_') | ||
} | ||
NAME_AND_TAGS_ARGS = { | ||
'display_name': _random_identifier('TestModel123_tags_'), | ||
'tags': ['test_tag123'] | ||
} | ||
FULL_MODEL_ARGS = { | ||
'display_name': _random_identifier('TestModel123_full_'), | ||
'tags': ['test_tag567'], | ||
'file_name': 'model1.tflite' | ||
} | ||
INVALID_FULL_MODEL_ARGS = { | ||
'display_name': _random_identifier('TestModel123_invalid_full_'), | ||
'tags': ['test_tag890'], | ||
'file_name': 'invalid_model.tflite' | ||
} | ||
|
||
|
||
@pytest.fixture | ||
def firebase_model(request): | ||
args = request.param | ||
tflite_format = None | ||
file_name = args.get('file_name') | ||
if file_name: | ||
file_path = testutils.resource_filename(file_name) | ||
source = ml.TFLiteGCSModelSource.from_tflite_model_file(file_path) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I assume each invocation of this makes an GCS upload? Do we need to clean them up at some point? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No. They overwrite and they are not randomized. So there should only ever be 3 files. |
||
tflite_format = ml.TFLiteFormat(model_source=source) | ||
|
||
ml_model = ml.Model( | ||
display_name=args.get('display_name'), | ||
tags=args.get('tags'), | ||
model_format=tflite_format) | ||
model = ml.create_model(model=ml_model) | ||
yield model | ||
_clean_up_model(model) | ||
|
||
|
||
@pytest.fixture | ||
def model_list(): | ||
ml_model_1 = ml.Model(display_name=_random_identifier('TestModel123_list1_')) | ||
model_1 = ml.create_model(model=ml_model_1) | ||
|
||
ml_model_2 = ml.Model(display_name=_random_identifier('TestModel123_list2_'), | ||
tags=['test_tag123']) | ||
model_2 = ml.create_model(model=ml_model_2) | ||
|
||
yield [model_1, model_2] | ||
|
||
_clean_up_model(model_1) | ||
_clean_up_model(model_2) | ||
|
||
|
||
def _clean_up_model(model): | ||
try: | ||
# Try to delete the model. | ||
# Some tests delete the model as part of the test. | ||
ml.delete_model(model.model_id) | ||
except exceptions.NotFoundError: | ||
pass | ||
|
||
|
||
# For rpc errors | ||
def check_firebase_error(excinfo, status, msg): | ||
err = excinfo.value | ||
assert isinstance(err, exceptions.FirebaseError) | ||
assert err.cause is not None | ||
assert err.http_response is not None | ||
assert err.http_response.status_code == status | ||
assert str(err) == msg | ||
|
||
|
||
# For operation errors | ||
def check_operation_error(excinfo, msg): | ||
err = excinfo.value | ||
assert isinstance(err, exceptions.FirebaseError) | ||
assert str(err) == msg | ||
|
||
|
||
def check_model(model, args): | ||
assert model.display_name == args.get('display_name') | ||
assert model.tags == args.get('tags') | ||
assert model.model_id is not None | ||
assert model.create_time is not None | ||
assert model.update_time is not None | ||
assert model.locked is False | ||
assert model.etag is not None | ||
|
||
|
||
def check_model_format(model, has_model_format=False, validation_error=None): | ||
if has_model_format: | ||
assert model.validation_error == validation_error | ||
assert model.published is False | ||
assert model.model_format.model_source.gcs_tflite_uri.startswith('gs://') | ||
if validation_error: | ||
assert model.model_format.size_bytes is None | ||
assert model.model_hash is None | ||
else: | ||
assert model.model_format.size_bytes is not None | ||
assert model.model_hash is not None | ||
else: | ||
assert model.model_format is None | ||
assert model.validation_error == 'No model file has been uploaded.' | ||
assert model.published is False | ||
assert model.model_hash is None | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_AND_TAGS_ARGS], indirect=True) | ||
def test_create_simple_model(firebase_model): | ||
check_model(firebase_model, NAME_AND_TAGS_ARGS) | ||
check_model_format(firebase_model) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [FULL_MODEL_ARGS], indirect=True) | ||
def test_create_full_model(firebase_model): | ||
check_model(firebase_model, FULL_MODEL_ARGS) | ||
check_model_format(firebase_model, True) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [FULL_MODEL_ARGS], indirect=True) | ||
def test_create_already_existing_fails(firebase_model): | ||
with pytest.raises(exceptions.AlreadyExistsError) as excinfo: | ||
ml.create_model(model=firebase_model) | ||
check_operation_error( | ||
excinfo, | ||
'Model \'{0}\' already exists'.format(firebase_model.display_name)) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [INVALID_FULL_MODEL_ARGS], indirect=True) | ||
def test_create_invalid_model(firebase_model): | ||
check_model(firebase_model, INVALID_FULL_MODEL_ARGS) | ||
check_model_format(firebase_model, True, 'Invalid flatbuffer format') | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_AND_TAGS_ARGS], indirect=True) | ||
def test_get_model(firebase_model): | ||
get_model = ml.get_model(firebase_model.model_id) | ||
check_model(get_model, NAME_AND_TAGS_ARGS) | ||
check_model_format(get_model) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_ONLY_ARGS], indirect=True) | ||
def test_get_non_existing_model(firebase_model): | ||
# Get a valid model_id that no longer exists | ||
ml.delete_model(firebase_model.model_id) | ||
|
||
with pytest.raises(exceptions.NotFoundError) as excinfo: | ||
ml.get_model(firebase_model.model_id) | ||
check_firebase_error(excinfo, 404, 'Requested entity was not found.') | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_ONLY_ARGS], indirect=True) | ||
def test_update_model(firebase_model): | ||
new_model_name = NAME_ONLY_ARGS_UPDATED.get('display_name') | ||
firebase_model.display_name = new_model_name | ||
updated_model = ml.update_model(firebase_model) | ||
check_model(updated_model, NAME_ONLY_ARGS_UPDATED) | ||
check_model_format(updated_model) | ||
|
||
# Second call with same model does not cause error | ||
updated_model2 = ml.update_model(updated_model) | ||
check_model(updated_model2, NAME_ONLY_ARGS_UPDATED) | ||
check_model_format(updated_model2) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_ONLY_ARGS], indirect=True) | ||
def test_update_non_existing_model(firebase_model): | ||
ml.delete_model(firebase_model.model_id) | ||
|
||
firebase_model.tags = ['tag987'] | ||
with pytest.raises(exceptions.NotFoundError) as excinfo: | ||
ml.update_model(firebase_model) | ||
check_operation_error( | ||
excinfo, | ||
'Model \'{0}\' was not found'.format(firebase_model.as_dict().get('name'))) | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [FULL_MODEL_ARGS], indirect=True) | ||
def test_publish_unpublish_model(firebase_model): | ||
assert firebase_model.published is False | ||
|
||
published_model = ml.publish_model(firebase_model.model_id) | ||
assert published_model.published is True | ||
|
||
unpublished_model = ml.unpublish_model(published_model.model_id) | ||
assert unpublished_model.published is False | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_ONLY_ARGS], indirect=True) | ||
def test_publish_invalid_fails(firebase_model): | ||
assert firebase_model.validation_error is not None | ||
|
||
with pytest.raises(exceptions.FailedPreconditionError) as excinfo: | ||
ml.publish_model(firebase_model.model_id) | ||
check_operation_error( | ||
excinfo, | ||
'Cannot publish a model that is not verified.') | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [FULL_MODEL_ARGS], indirect=True) | ||
def test_publish_unpublish_non_existing_model(firebase_model): | ||
ml.delete_model(firebase_model.model_id) | ||
|
||
with pytest.raises(exceptions.NotFoundError) as excinfo: | ||
ml.publish_model(firebase_model.model_id) | ||
check_operation_error( | ||
excinfo, | ||
'Model \'{0}\' was not found'.format(firebase_model.as_dict().get('name'))) | ||
|
||
with pytest.raises(exceptions.NotFoundError) as excinfo: | ||
ml.unpublish_model(firebase_model.model_id) | ||
check_operation_error( | ||
excinfo, | ||
'Model \'{0}\' was not found'.format(firebase_model.as_dict().get('name'))) | ||
|
||
|
||
def test_list_models(model_list): | ||
filter_str = 'displayName={0} OR tags:{1}'.format( | ||
model_list[0].display_name, model_list[1].tags[0]) | ||
|
||
all_models = ml.list_models(list_filter=filter_str) | ||
all_model_ids = [mdl.model_id for mdl in all_models.iterate_all()] | ||
for mdl in model_list: | ||
assert mdl.model_id in all_model_ids | ||
|
||
|
||
def test_list_models_invalid_filter(): | ||
invalid_filter = 'InvalidFilterParam=123' | ||
|
||
with pytest.raises(exceptions.InvalidArgumentError) as excinfo: | ||
ml.list_models(list_filter=invalid_filter) | ||
check_firebase_error(excinfo, 400, 'Request contains an invalid argument.') | ||
|
||
|
||
@pytest.mark.parametrize('firebase_model', [NAME_ONLY_ARGS], indirect=True) | ||
def test_delete_model(firebase_model): | ||
ml.delete_model(firebase_model.model_id) | ||
|
||
# Second delete of same model will fail | ||
with pytest.raises(exceptions.NotFoundError) as excinfo: | ||
ml.delete_model(firebase_model.model_id) | ||
check_firebase_error(excinfo, 404, 'Requested entity was not found.') | ||
|
||
|
||
# Test tensor flow conversion functions if tensor flow is enabled. | ||
#'pip install tensorflow' in the environment if you want _TF_ENABLED = True | ||
#'pip install tensorflow==2.0.0b' for version 2 etc. | ||
|
||
|
||
def _clean_up_directory(save_dir): | ||
if save_dir.startswith(tempfile.gettempdir()) and os.path.exists(save_dir): | ||
shutil.rmtree(save_dir) | ||
|
||
|
||
@pytest.fixture | ||
def keras_model(): | ||
assert _TF_ENABLED | ||
x_array = [-1, 0, 1, 2, 3, 4] | ||
y_array = [-3, -1, 1, 3, 5, 7] | ||
model = tf.keras.models.Sequential( | ||
[tf.keras.layers.Dense(units=1, input_shape=[1])]) | ||
model.compile(optimizer='sgd', loss='mean_squared_error') | ||
model.fit(x_array, y_array, epochs=3) | ||
return model | ||
|
||
|
||
@pytest.fixture | ||
def saved_model_dir(keras_model): | ||
assert _TF_ENABLED | ||
# Make a new parent directory. The child directory must not exist yet. | ||
# The child directory gets created by tf. If it exists, the tf call fails. | ||
parent = tempfile.mkdtemp() | ||
save_dir = os.path.join(parent, 'child') | ||
|
||
# different versions have different model conversion capability | ||
# pick something that works for each version | ||
if tf.version.VERSION.startswith('1.'): | ||
tf.reset_default_graph() | ||
x_var = tf.placeholder(tf.float32, (None, 3), name="x") | ||
y_var = tf.multiply(x_var, x_var, name="y") | ||
with tf.Session() as sess: | ||
tf.saved_model.simple_save(sess, save_dir, {"x": x_var}, {"y": y_var}) | ||
else: | ||
# If it's not version 1.x or version 2.x we need to update the test. | ||
assert tf.version.VERSION.startswith('2.') | ||
tf.saved_model.save(keras_model, save_dir) | ||
yield save_dir | ||
_clean_up_directory(parent) | ||
|
||
|
||
@pytest.mark.skipif(not _TF_ENABLED, reason='Tensor flow is required for this test.') | ||
def test_from_keras_model(keras_model): | ||
source = ml.TFLiteGCSModelSource.from_keras_model(keras_model, 'model2.tflite') | ||
assert re.search( | ||
'^gs://.*/Firebase/ML/Models/model2.tflite$', | ||
source.gcs_tflite_uri) is not None | ||
|
||
# Validate the conversion by creating a model | ||
model_format = ml.TFLiteFormat(model_source=source) | ||
model = ml.Model(display_name=_random_identifier('KerasModel_'), model_format=model_format) | ||
created_model = ml.create_model(model) | ||
|
||
try: | ||
check_model(created_model, {'display_name': model.display_name}) | ||
check_model_format(created_model, True) | ||
finally: | ||
_clean_up_model(created_model) | ||
|
||
|
||
@pytest.mark.skipif(not _TF_ENABLED, reason='Tensor flow is required for this test.') | ||
def test_from_saved_model(saved_model_dir): | ||
# Test the conversion helper | ||
source = ml.TFLiteGCSModelSource.from_saved_model(saved_model_dir, 'model3.tflite') | ||
assert re.search( | ||
'^gs://.*/Firebase/ML/Models/model3.tflite$', | ||
source.gcs_tflite_uri) is not None | ||
|
||
# Validate the conversion by creating a model | ||
model_format = ml.TFLiteFormat(model_source=source) | ||
model = ml.Model(display_name=_random_identifier('SavedModel_'), model_format=model_format) | ||
created_model = ml.create_model(model) | ||
|
||
try: | ||
assert created_model.model_id is not None | ||
assert created_model.validation_error is None | ||
finally: | ||
_clean_up_model(created_model) |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
This is not a tflite file. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
for _ in range(8)
and remove the ignore directive.