From b1d861967543f4384663d4994b793e81aacd9c89 Mon Sep 17 00:00:00 2001
From: Augustin Zidek <augustinzidek@google.com>
Date: Fri, 1 Aug 2025 04:18:25 -0700
Subject: [PATCH] Use pytest for testing instead of setup.py test

PiperOrigin-RevId: 789706772
Change-Id: Id4d58250e03bc41ab7d9a90c3d569266cb2db4ad
---
 alphafold/notebooks/notebook_utils_test.py |  2 +-
 conftest.py                                | 30 ++++++++++++++++++++++
 requirements.txt                           |  4 ++-
 run_alphafold_test.py                      |  4 +--
 setup.py                                   |  6 ++---
 5 files changed, 37 insertions(+), 9 deletions(-)
 create mode 100644 conftest.py

diff --git a/alphafold/notebooks/notebook_utils_test.py b/alphafold/notebooks/notebook_utils_test.py
index a67ceec..6df7689 100644
--- a/alphafold/notebooks/notebook_utils_test.py
+++ b/alphafold/notebooks/notebook_utils_test.py
@@ -14,6 +14,7 @@
 
 """Tests for notebook_utils."""
 import io
+from unittest import mock
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -21,7 +22,6 @@ from alphafold.data import parsers
 from alphafold.data import templates
 from alphafold.notebooks import notebook_utils
 
-import mock
 import numpy as np
 
 
diff --git a/conftest.py b/conftest.py
new file mode 100644
index 0000000..5f2821a
--- /dev/null
+++ b/conftest.py
@@ -0,0 +1,30 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Fixture for pytest.
+
+This is needed to parse the absl flags before running the test.
+"""
+
+import sys
+
+from absl import flags
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def initialize_absl_flags(request):
+  del request
+  # Parse any flags that make sense to absl as absl flags.
+  flags.FLAGS(sys.argv, known_only=True)
diff --git a/requirements.txt b/requirements.txt
index ed442fa..fce966b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,9 +6,11 @@ dm-tree==0.1.8
 docker==5.0.0
 immutabledict==2.0.0
 jax==0.4.26
+matplotlib==3.8.0
 ml-collections==0.1.0
 numpy==1.24.3
 pandas==2.0.3
-setuptools<72.0.0
+pytest<8.5.0
 scipy==1.11.1
+setuptools<72.0.0
 tensorflow-cpu==2.16.1
diff --git a/run_alphafold_test.py b/run_alphafold_test.py
index f9da3a0..42c1a79 100644
--- a/run_alphafold_test.py
+++ b/run_alphafold_test.py
@@ -12,15 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for run_alphafold."""
-
 import json
 import os
+from unittest import mock
 
 from absl.testing import absltest
 from absl.testing import parameterized
 import run_alphafold
-import mock
 import numpy as np
 # Internal import (7716).
 
diff --git a/setup.py b/setup.py
index 2252af6..bfa18df 100644
--- a/setup.py
+++ b/setup.py
@@ -39,16 +39,14 @@ setup(
         'docker',
         'immutabledict',
         'jax',
+        'matplotlib',
         'ml-collections',
         'numpy',
         'pandas',
+        'pytest',
         'scipy',
         'tensorflow-cpu',
     ],
-    tests_require=[
-        'matplotlib',  # For notebook_utils_test.
-        'mock',
-    ],
     classifiers=[
         'Development Status :: 5 - Production/Stable',
         'Intended Audience :: Science/Research',