Merge pull request #31 from Yoctol/more_test_cases

case: save more than one time

Author: SoluMilken

serving_utils/tests/test_saver.py

@@ -75,31 +75,36 @@ def test_save_correct_files(saver, output_dir):
     assert expected_output_filenames == real_output_filenames
 
 
+def test_save_several_times(saver, output_dir):
+    n = 5
+    for _ in range(n):
+        saver.save()
+
+    expected_output_filenames = set(
+        [
+            join(output_dir, f"{i}", filename)
+            for i in range(n) for filename in [
+                'saved_model.pb',
+                'variables',
+                'variables/variables.data-00000-of-00001',
+                'variables/variables.index',
+            ]
+        ],
+    )
+    real_output_filenames = set(
+        glob(output_dir + '/*/*') +  # noqa: W504
+        glob(output_dir + '/*/*/*'),
+    )
+    assert expected_output_filenames == real_output_filenames
+
+
 def test_save_will_not_change_model(x_test, y_test, model, saver, output_dir):
     old_loss = model.evaluate(x_test, y_test)
     saver.save()
-    new_loss = load_n_evaluate(model.sess, output_dir, x_test, y_test)
+    new_loss = load_n_evaluate(output_dir, x_test, y_test)
     assert old_loss == new_loss
 
 
-def load_n_evaluate(sess, output_dir, x_test, y_test):
-    with tf.Session(graph=tf.Graph()) as sess:
-        meta_graph_def = tf.saved_model.loader.load(
-            sess=sess,
-            tags=[tf.saved_model.tag_constants.SERVING],
-            export_dir=join(output_dir, '0'),
-        )
-        evaluate_graph = meta_graph_def.signature_def['evaluate']
-        loss = sess.run(
-            evaluate_graph.outputs['loss'].name,
-            feed_dict={
-                evaluate_graph.inputs['x'].name: x_test,
-                evaluate_graph.inputs['y'].name: y_test,
-            },
-        )
-    return loss
-
-
 def test_frozen_save_correct_files(saver_with_freeze, output_dir):
     saver_with_freeze.save()
     expected_output_filenames = set(
@@ -116,19 +121,48 @@ def test_frozen_save_correct_files(saver_with_freeze, output_dir):
     )
     assert expected_output_filenames == real_output_filenames
 
-    # def test_freeze_graph_has_session_update(self):
-    #     old_sess = saver.session
-    #     saver.freeze_graph()
-    #     new_sess = saver.session
-    #     assertNotEqual(old_sess, new_sess)
-
-    # def test_freeze_graph_will_not_change_loss(self):
-    #     old_loss = model.evaluate(x_test, y_test)
-    #     saver.freeze_graph()
-    #     model.sess = saver.session
-    #     new_loss = model.evaluate(x_test, y_test)
-    #     assertEqual(old_loss, new_loss)
-
-    # def test_freeze_n_save(self):
-    #     saver.freeze_graph()
-    #     saver.save()
+
+def test_frozen_save_several_times(saver_with_freeze, output_dir):
+    n = 5
+    for _ in range(n):
+        saver_with_freeze.save()
+
+    expected_output_filenames = set(
+        [
+            join(output_dir, f"{i}", filename)
+            for i in range(n) for filename in [
+                'saved_model.pb',
+                'variables',
+            ]
+        ],
+    )
+    real_output_filenames = set(
+        glob(output_dir + '/*/*') +  # noqa: W504
+        glob(output_dir + '/*/*/*'),
+    )
+    assert expected_output_filenames == real_output_filenames
+
+
+def test_frozen_save_will_not_change_model(x_test, y_test, model, saver_with_freeze, output_dir):
+    old_loss = model.evaluate(x_test, y_test)
+    saver_with_freeze.save()
+    new_loss = load_n_evaluate(output_dir, x_test, y_test)
+    assert old_loss == new_loss
+
+
+def load_n_evaluate(output_dir, x_test, y_test):
+    with tf.Session(graph=tf.Graph()) as sess:
+        meta_graph_def = tf.saved_model.loader.load(
+            sess=sess,
+            tags=[tf.saved_model.tag_constants.SERVING],
+            export_dir=join(output_dir, '0'),
+        )
+        evaluate_graph = meta_graph_def.signature_def['evaluate']
+        loss = sess.run(
+            evaluate_graph.outputs['loss'].name,
+            feed_dict={
+                evaluate_graph.inputs['x'].name: x_test,
+                evaluate_graph.inputs['y'].name: y_test,
+            },
+        )
+    return loss