diff --git a/test.py.in b/test.py.in
index 7f1418a9d0..4221e556d8 100644
--- a/test.py.in
+++ b/test.py.in
@@ -4,7 +4,9 @@ from ortools.linear_solver import linear_solver_pb2
 from ortools.constraint_solver import pywrapcp
 from ortools.sat.python import swig_helper
 from ortools.sat.python import cp_model
-from ortools.graph import pywrapgraph
+from ortools.graph.python import linear_sum_assignment
+from ortools.graph.python import max_flow
+from ortools.graph.python import min_cost_flow
 from ortools.algorithms import pywrapknapsack_solver
 from ortools.scheduling import pywraprcpsp
 from ortools.util.python import sorted_interval_list