Backport mock assertions for Python 2 tests

tensorboard/tools/whitespace_hygiene_test.py

@@ -28,8 +28,10 @@
 import sys
 
 
-# Remove files from this list as whitespace errors are fixed.
 exceptions = frozenset([
+    # End-of-line whitespace is semantic in patch files when a line
+    # contains a single space.
+    "third_party/mock_call_assertions.patch",
 ])
 
 

tensorboard/uploader/exporter_test.py

@@ -27,11 +27,9 @@
 
 try:
   # python version >= 3.3
-  from unittest import mock
+  from unittest import mock  # pylint: disable=g-import-not-at-top
 except ImportError:
-  # mock==1.0.0 backport lacks `assert_called_once` and friends
-  print("Test disabled in Python 2")
-  exit(0)
+  import mock  # pylint: disable=g-import-not-at-top,unused-import
 
 
 from tensorboard.uploader.proto import export_service_pb2

tensorboard/uploader/uploader_test.py

@@ -26,11 +26,9 @@
 
 try:
   # python version >= 3.3
-  from unittest import mock
+  from unittest import mock  # pylint: disable=g-import-not-at-top
 except ImportError:
-  # mock==1.0.0 backport lacks `assert_called_once` and friends
-  print("Test disabled in Python 2")
-  exit(0)
+  import mock  # pylint: disable=g-import-not-at-top,unused-import
 
 import tensorflow as tf
 

tensorboard/uploader/util_test.py

@@ -24,11 +24,9 @@
 
 try:
   # python version >= 3.3
-  from unittest import mock
+  from unittest import mock  # pylint: disable=g-import-not-at-top
 except ImportError:
-  # mock==1.0.0 backport lacks `assert_called_once` and friends
-  print("Test disabled in Python 2")
-  exit(0)
+  import mock  # pylint: disable=g-import-not-at-top,unused-import
 
 
 from google.protobuf import timestamp_pb2

third_party/mock_call_assertions.patch

@@ -0,0 +1,59 @@
+--- mock.py	2012-10-07 18:00:10.000000000 +0100
++++ mock.py	2019-10-24 22:19:25.657417082 -0700
+@@ -286,6 +286,12 @@
+     if not _is_instance_mock(mock):
+         return
+
++    def assert_called(*args, **kwargs):
++        return mock.assert_called(*args, **kwargs)
++    def assert_not_called(*args, **kwargs):
++        return mock.assert_not_called(*args, **kwargs)
++    def assert_called_once(*args, **kwargs):
++        return mock.assert_called_once(*args, **kwargs)
+     def assert_called_with(*args, **kwargs):
+         return mock.assert_called_with(*args, **kwargs)
+     def assert_called_once_with(*args, **kwargs):
+@@ -318,6 +324,9 @@
+     funcopy.assert_has_calls = assert_has_calls
+     funcopy.assert_any_call = assert_any_call
+     funcopy.reset_mock = reset_mock
++    funcopy.assert_called = assert_called
++    funcopy.assert_not_called = assert_not_called
++    funcopy.assert_called_once = assert_called_once
+
+     mock._mock_delegate = funcopy
+
+@@ -809,6 +818,33 @@
+         return message % (expected_string, actual_string)
+
+
++    def assert_not_called(_mock_self):
++        """assert that the mock was never called.
++        """
++        self = _mock_self
++        if self.call_count != 0:
++            msg = ("Expected '%s' to not have been called. Called %s times." %
++                   (self._mock_name or 'mock', self.call_count))
++            raise AssertionError(msg)
++
++    def assert_called(_mock_self):
++        """assert that the mock was called at least once
++        """
++        self = _mock_self
++        if self.call_count == 0:
++            msg = ("Expected '%s' to have been called." %
++                   self._mock_name or 'mock')
++            raise AssertionError(msg)
++
++    def assert_called_once(_mock_self):
++        """assert that the mock was called only once.
++        """
++        self = _mock_self
++        if not self.call_count == 1:
++            msg = ("Expected '%s' to have been called once. Called %s times." %
++                   (self._mock_name or 'mock', self.call_count))
++            raise AssertionError(msg)
++
+     def assert_called_with(_mock_self, *args, **kwargs):
+         """assert that the mock was called with the specified arguments.
+

third_party/python.bzl

@@ -114,6 +114,14 @@ def tensorboard_python_workspace():
         sha256 = "2d9fbe67001d2e8f02692075257f3c11e1b0194bd838c8ce3f49b31fc6c3f033",
         strip_prefix = "mock-1.0.0",
         build_file = str(Label("//third_party:mock.BUILD")),
+        patches = [
+            # `mock==1.0.0` lacks some assertion methods present in
+            # later versions of `mock` (see comment above for why we pin
+            # to this version). Patch created by diffing the pinned
+            # `mock.py` with GitHub head and identifying all the bits
+            # that looked related to the methods in question.
+            "//third_party:mock_call_assertions.patch",
+        ],
     )
 
     http_archive(