def test_switch_back_modified_path(self):
new_path = os.pathsep.join(("/bar/foo", "quux/baz"))
os.environ["PATH"] = os.pathsep.join(
(new_path, self._original_os_path))
os.environ["PYTHONPATH"] = os.pathsep.join(
(new_path, self._original_os_pythonpath))
_tensorflow_magics._tensorflow_version("2.x")
_tensorflow_magics._tensorflow_version("1.x")
tf1_path = _tensorflow_magics._VERSIONS["1"].path
self._assert_starts_with(sys.path[0], tf1_path)
self.assertEqual(sys.path, self._original_sys_path)
(path_head, path_tail) = os.environ["PATH"].split(os.pathsep, 1)
self._assert_starts_with(path_head, tf1_path)
self._assert_ends_with(path_head, "bin")
self.assertEqual(
path_tail,
os.pathsep.join([new_path] +
self._original_os_path.split(os.pathsep, 1)[1:]))
(pythonpath_head,
pythonpath_tail) = os.environ["PYTHONPATH"].split(os.pathsep, 1)
self._assert_starts_with(pythonpath_head, tf1_path)
self.assertEqual(
pythonpath_tail,
os.pathsep.join(
[new_path] +
self._original_os_pythonpath.split(os.pathsep, 1)[1:]))
def test_switch_back_default_path(self):
_tensorflow_magics._tensorflow_version("2.x")
_tensorflow_magics._tensorflow_version("1.x")
self.assertEqual(sys.path, self._original_sys_path)
self.assertEqual(os.environ["PATH"], self._original_os_path)
self.assertEqual(os.environ["PYTHONPATH"],
self._original_os_pythonpath)
def test_switch_back_no_paths(self):
os.environ.pop("PATH", None)
os.environ.pop("PYTHONPATH", None)
_tensorflow_magics._tensorflow_version("2.x")
_tensorflow_magics._tensorflow_version("1.x")
self.assertEqual(sys.path, self._original_sys_path)
self.assertEqual(os.environ.get("PATH", ""), "")
self.assertEqual(os.environ.get("PYTHONPATH", ""), "")
def test_switch_1x_to_2x_default_path(self):
_tensorflow_magics._tensorflow_version("2.x")
self.assertEqual(sys.path, self._original_sys_path[1:])
self.assertEqual(
os.environ["PATH"],
os.pathsep.join(self._original_os_path.split(os.pathsep, 1)[1:]))
self.assertEqual(
os.environ["PYTHONPATH"],
os.pathsep.join(
self._original_os_pythonpath.split(os.pathsep, 1)[1:]))
def test_switch_back_with_paths(self):
original_pythonpath = os.pathsep.join(("/foo/bar", "/baz/quux"))
original_os_path = os.pathsep.join(("/bar/foo", "/quux/baz"))
os.environ["PYTHONPATH"] = original_pythonpath
os.environ["PATH"] = original_os_path
_tensorflow_magics._tensorflow_version("1.x")
_tensorflow_magics._tensorflow_version("2.x")
self.assertEqual(sys.path, self._original_sys_path)
self.assertEqual(os.environ["PATH"], original_os_path)
self.assertEqual(os.environ["PYTHONPATH"], original_pythonpath)
def test_tpu_version_switch(self):
_tensorflow_magics._get_tf_version = mock.Mock(
return_value=_tensorflow_magics._VERSIONS["2"].version)
_tensorflow_magics._tensorflow_version("2.x")
_tensorflow_magics._get_tf_version = mock.Mock(
return_value=_tensorflow_magics._VERSIONS["1"].version)
_tensorflow_magics._tensorflow_version("1.x")
expected = "http://0.0.0.0:8475/requestversion/{}"
calls = [
mock.call(expected.format(_tensorflow_magics._VERSIONS["2"].version)),
mock.call(expected.format(_tensorflow_magics._VERSIONS["1"].version)),
]
self.assertEqual(requests.post.mock_calls, calls)
def test_handle_tf_install_after_setting_version(self):
_tensorflow_magics._tensorflow_version("1.x")
_tensorflow_magics._handle_tf_install()
# _handle_tf_install should be a no-op because magic was invoked.
self.assertTrue(_tensorflow_magics._explicitly_set())
tf1_path = _tensorflow_magics._VERSIONS["1"].path
self._assert_starts_with(sys.path[0], tf1_path)
path_head = os.environ["PATH"].split(os.pathsep)[0]
self._assert_starts_with(path_head, tf1_path)
self._assert_ends_with(path_head, "bin")
self._assert_starts_with(os.environ["PYTHONPATH"].split(os.pathsep)[0],
tf1_path)
def test_switch_1x_to_2x_no_paths(self):
os.environ.pop("PATH", None)
os.environ.pop("PYTHONPATH", None)
tf2_path = _tensorflow_magics._VERSIONS["2"].path
_tensorflow_magics._tensorflow_version("2.x")
self.assertEqual(sys.path[1:], self._original_sys_path)
self._assert_starts_with(sys.path[0], tf2_path)
self._assert_starts_with(os.environ["PYTHONPATH"], tf2_path)
self._assert_len(os.environ["PYTHONPATH"].split(os.pathsep), 1)
os_path_head, os_path_tail = os.environ["PATH"].split(os.pathsep, 1)
self._assert_starts_with(os_path_head, tf2_path)
self._assert_ends_with(os_path_head, "bin")
self.assertEqual(os_path_tail, "")
def test_switch_1x_to_2x_modified_path(self):
new_path = os.pathsep.join(("/bar/foo", "quux/baz"))
os.environ["PATH"] = os.pathsep.join(
(new_path, self._original_os_path))
os.environ["PYTHONPATH"] = os.pathsep.join(
(new_path, self._original_os_pythonpath))
_tensorflow_magics._tensorflow_version("2.x")
self.assertEqual(sys.path, self._original_sys_path[1:])
self.assertEqual(
os.environ["PATH"],
os.pathsep.join([new_path] +
self._original_os_path.split(os.pathsep, 1)[1:]))
self.assertEqual(
os.environ["PYTHONPATH"],
os.pathsep.join(
[new_path] +
self._original_os_pythonpath.split(os.pathsep, 1)[1:]))
def test_switch_1x_to_2x_existing_paths(self):
original_pythonpath = os.pathsep.join(("/foo/bar", "/baz/quux"))
original_os_path = os.pathsep.join(("/bar/foo", "/quux/baz"))
os.environ["PYTHONPATH"] = original_pythonpath
os.environ["PATH"] = original_os_path
tf2_path = _tensorflow_magics._VERSIONS["2"].path
_tensorflow_magics._tensorflow_version("2.x")
self.assertEqual(sys.path[1:], self._original_sys_path)
self._assert_starts_with(sys.path[0], tf2_path)
(pythonpath_head,
pythonpath_tail) = os.environ["PYTHONPATH"].split(os.pathsep, 1)
self._assert_starts_with(pythonpath_head, tf2_path)
self.assertEqual(pythonpath_tail, original_pythonpath)
(os_path_head, os_path_tail) = os.environ["PATH"].split(os.pathsep, 1)
self._assert_starts_with(os_path_head, tf2_path)
self._assert_ends_with(os_path_head, "bin")
self.assertEqual(os_path_tail, original_os_path)
def test_tpu_version_switch(self):
_tensorflow_magics._get_tf_version = mock.Mock(
return_value=_tensorflow_magics._VERSIONS["1"].version)
_tensorflow_magics._tensorflow_version("1.x")
_tensorflow_magics._get_tf_version = mock.Mock(
return_value=_tensorflow_magics._VERSIONS["2"].version)
_tensorflow_magics._tensorflow_version("2.x")
expected = "http://0.0.0.0:8475/requestversion/{}"
# TODO(b/151765674): Remove this comment and initial call once TPU defaults
# to TF1.
# We expect one additional call from initialization as the TPU VM defaults
# to TF1 and we default to TF2.
calls = [
mock.call(
expected.format(_tensorflow_magics._VERSIONS["2"].version)),
mock.call(
expected.format(_tensorflow_magics._VERSIONS["1"].version)),
mock.call(
expected.format(_tensorflow_magics._VERSIONS["2"].version)),
]
self.assertEqual(requests.post.mock_calls, calls)
def test_switch_back_does_not_import(self):
_tensorflow_magics._tensorflow_version("2.x")
_tensorflow_magics._tensorflow_version("1.x")
self.assertNotIn("tensorflow", sys.modules)
