@@ -27,11 +27,18 @@ def test_copy_to_with_psycopg2(self):
27
27
"""Test copy_to function with psycopg2."""
28
28
# Mock the psycopg2 import to succeed and psycopg to fail
29
29
with mock .patch .dict ("sys.modules" , {"psycopg" : None }):
30
+ # Reload the module to ensure it uses the psycopg2 implementation
31
+ import importlib
32
+ import postgres_copy .psycopg_compat
33
+
34
+ importlib .reload (postgres_copy .psycopg_compat )
35
+ from postgres_copy .psycopg_compat import copy_to as reloaded_copy_to
36
+
30
37
with mock .patch ("psycopg2.extensions.adapt" ) as mock_adapt :
31
38
mock_adapt .return_value .getquoted .return_value = b"'test'"
32
39
33
40
# Call the function
34
- copy_to (self .cursor , self .sql , (1 , 2 ), self .destination )
41
+ reloaded_copy_to (self .cursor , self .sql , (1 , 2 ), self .destination )
35
42
36
43
# Check that the psycopg2 version was called with the right parameters
37
44
self .cursor .copy_expert .assert_called_once ()
@@ -45,8 +52,15 @@ def test_copy_from_with_psycopg2(self):
45
52
"""Test copy_from function with psycopg2."""
46
53
# Mock the psycopg2 import to succeed and psycopg to fail
47
54
with mock .patch .dict ("sys.modules" , {"psycopg" : None }):
55
+ # Reload the module to ensure it uses the psycopg2 implementation
56
+ import importlib
57
+ import postgres_copy .psycopg_compat
58
+
59
+ importlib .reload (postgres_copy .psycopg_compat )
60
+ from postgres_copy .psycopg_compat import copy_from as reloaded_copy_from
61
+
48
62
# Call the function
49
- copy_from (self .cursor , self .sql , self .source )
63
+ reloaded_copy_from (self .cursor , self .sql , self .source )
50
64
51
65
# Check that the psycopg2 version was called with the right parameters
52
66
self .cursor .copy_expert .assert_called_once ()
0 commit comments