diff --git a/grain/_src/python/data_sources_test.py b/grain/_src/python/data_sources_test.py
index d8007d74a..bdcd9834e 100644
--- a/grain/_src/python/data_sources_test.py
+++ b/grain/_src/python/data_sources_test.py
@@ -14,6 +14,7 @@
 """Tests for data sources."""
 
 from collections.abc import Sequence
+import copy
 import dataclasses
 import pathlib
 import pickle