@@ -431,8 +431,8 @@ def test__get_foreign_keys(self):
431431 def test__get_primary_keys_info (
432432 self , _diff_dicts_mock , _get_primary_keys_mock ):
433433 _get_primary_keys_mock .side_effect = [
434- ['pk_left_1' , 'pk_left_2' ],
435- ['pk_right_1' ]
434+ { 'constrained_columns' : ['pk_left_1' , 'pk_left_2' ]} ,
435+ { 'constrained_columns' : ['pk_right_1' ]}
436436 ]
437437 left_inspector , right_inspector = Mock (), Mock ()
438438
@@ -449,8 +449,8 @@ def test__get_primary_keys_info(
449449 def test__get_primary_keys_info_ignores (
450450 self , _diff_dicts_mock , _get_primary_keys_mock ):
451451 _get_primary_keys_mock .side_effect = [
452- ['pk_left_1' , 'pk_left_2' ],
453- ['pk_right_1' , 'pk_right_2' ]
452+ { 'constrained_columns' : ['pk_left_1' , 'pk_left_2' ]} ,
453+ { 'constrained_columns' : ['pk_right_1' , 'pk_right_2' ]},
454454 ]
455455 left_inspector , right_inspector = Mock (), Mock ()
456456 ignores = ['pk_left_1' , 'pk_right_2' ]
@@ -465,13 +465,58 @@ def test__get_primary_keys_info_ignores(
465465
466466 assert _diff_dicts_mock .return_value == result
467467
468+ def test__get_primary_keys_info_with_pk_constraint_name (
469+ self , _diff_dicts_mock , _get_primary_keys_mock ):
470+ _get_primary_keys_mock .side_effect = [
471+ {'name' : 'left' , 'constrained_columns' : ['pk_left_1' ]},
472+ {'name' : 'right' , 'constrained_columns' : ['pk_right_1' ]}
473+ ]
474+ left_inspector , right_inspector = Mock (), Mock ()
475+
476+ result = _get_primary_keys_info (
477+ left_inspector , right_inspector , 'table_A' , [])
478+
479+ _diff_dicts_mock .assert_called_once_with (
480+ {
481+ 'left' : {'name' : 'left' ,
482+ 'constrained_columns' : ['pk_left_1' ]}
483+ },
484+ {
485+ 'right' : {'name' : 'right' ,
486+ 'constrained_columns' : ['pk_right_1' ]}
487+ }
488+ )
489+ assert _diff_dicts_mock .return_value == result
490+
491+ def test__get_primary_keys_info_ignores_with_pk_constraint_name (
492+ self , _diff_dicts_mock , _get_primary_keys_mock ):
493+ _get_primary_keys_mock .side_effect = [
494+ {'name' : 'left_1' , 'constrained_columns' : ['pk_left_1' ]},
495+ {'name' : 'right_1' , 'constrained_columns' : ['pk_right_1' ]},
496+ ]
497+ left_inspector , right_inspector = Mock (), Mock ()
498+ ignores = ['left_1' , 'left_2' , 'right_2' ]
499+
500+ result = _get_primary_keys_info (
501+ left_inspector , right_inspector , 'table_A' , ignores )
502+
503+ _diff_dicts_mock .assert_called_once_with (
504+ dict (),
505+ {
506+ 'right_1' : {'name' : 'right_1' ,
507+ 'constrained_columns' : ['pk_right_1' ]},
508+ }
509+ )
510+
511+ assert _diff_dicts_mock .return_value == result
512+
468513 def test__get_primary_keys (self ):
469514 inspector = Mock ()
470515
471516 result = _get_primary_keys (inspector , 'table_A' )
472517
473- inspector .get_primary_keys .assert_called_once_with ('table_A' )
474- assert inspector .get_primary_keys .return_value == result
518+ inspector .get_pk_constraint .assert_called_once_with ('table_A' )
519+ assert inspector .get_pk_constraint .return_value == result
475520
476521 def test__get_indexes_info (
477522 self , _diff_dicts_mock , _get_indexes_mock ):
0 commit comments