Move from deprecated self.test_session() to self.session() when a graph is set.

self.test_session() has been deprecated in cl/208545396 as its behavior confuses readers of the test. Moving to self.session() instead.

PiperOrigin-RevId: 209696110

2 files changed, 8 insertions, 8 deletions

diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
index 63b9c5fddf..906c916b27 100644
--- a/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
+++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/model_ops_test.py
@@ -133,7 +133,7 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
 
   def testSerialization(self):
     with ops.Graph().as_default() as graph:
-      with self.test_session(graph):
+      with self.session(graph):
         tree_ensemble_config = tree_config_pb2.DecisionTreeEnsembleConfig()
         # Bias tree only for second class.
         tree1 = tree_ensemble_config.trees.add()
@@ -164,7 +164,7 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
         serialized_config = serialized_config.eval()
 
     with ops.Graph().as_default() as graph:
-      with self.test_session(graph):
+      with self.session(graph):
         tree_ensemble_handle2 = model_ops.tree_ensemble_variable(
             stamp_token=9,
             tree_ensemble_config=serialized_config,
@@ -211,7 +211,7 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
     # to ensure no caching occurs under the hood.
     save_path = os.path.join(self.get_temp_dir(), "restore-test")
     with ops.Graph().as_default() as graph:
-      with self.test_session(graph) as sess:
+      with self.session(graph) as sess:
         # Prepare learner config.
         learner_config = learner_pb2.LearnerConfig()
         learner_config.num_classes = 2
@@ -288,7 +288,7 @@ class ModelOpsTest(test_util.TensorFlowTestCase):
     # Start a second session.  In that session the parameter nodes
     # have not been initialized either.
     with ops.Graph().as_default() as graph:
-      with self.test_session(graph) as sess:
+      with self.session(graph) as sess:
         tree_ensemble_handle = model_ops.tree_ensemble_variable(
             stamp_token=0, tree_ensemble_config="", name="restore_tree")
         my_saver = saver.Saver()
diff --git a/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py
index 074623699d..cddb159f82 100644
--- a/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py
+++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/quantile_ops_test.py
@@ -338,7 +338,7 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase):
     save_dir = os.path.join(self.get_temp_dir(), "save_restore")
     save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
 
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       accumulator = quantile_ops.QuantileAccumulator(
           init_stamp_token=0, num_quantiles=3, epsilon=0.33, name="q0")
 
@@ -366,7 +366,7 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase):
       self.assertEqual(True, are_ready_flush)
       self.assertAllEqual([2, 4, 6.], buckets)
 
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       accumulator = quantile_ops.QuantileAccumulator(
           init_stamp_token=0, num_quantiles=3, epsilon=0.33, name="q0")
       save = saver.Saver()
@@ -389,7 +389,7 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase):
     save_dir = os.path.join(self.get_temp_dir(), "save_restore")
     save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), "hash")
 
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       accumulator = quantile_ops.QuantileAccumulator(
           init_stamp_token=0, num_quantiles=3, epsilon=0.33, name="q0")
 
@@ -413,7 +413,7 @@ class QuantileBucketsOpTest(test_util.TensorFlowTestCase):
       self.assertAllEqual([1, 3, 5], buckets)
       save.save(sess, save_path)
 
-    with self.test_session(graph=ops.Graph()) as sess:
+    with self.session(graph=ops.Graph()) as sess:
       accumulator = quantile_ops.QuantileAccumulator(
           init_stamp_token=0, num_quantiles=3, epsilon=0.33, name="q0")
       save = saver.Saver()