Fixing a unittest

PiperOrigin-RevId: 743221887

Author: lukebaumann

pathwaysutils/test/proxy_backend_test.py

@@ -11,10 +11,13 @@
 # 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.
+"""Tests for the proxy backend module."""
+
+from unittest import mock
+
 import jax
 from jax.extend import backend
 from jax.lib.xla_extension import ifrt_proxy
-import mock
 from pathwaysutils import proxy_backend
 
 from absl.testing import absltest