test: fix expected ranges

larkee · larkee · commit 50168d2307ca · 2020-11-03T16:15:55.000+11:00
diff --git a/tests/unit/test_keyset.py b/tests/unit/test_keyset.py
@@ -303,6 +303,7 @@ def test_to_pb_w_only_keys(self):
         self.assertEqual(len(result.ranges), 0)
 
     def test_to_pb_w_only_ranges(self):
+        from google.cloud.spanner_v1 import KeyRangePB
         from google.cloud.spanner_v1 import KeySetPB
         from google.cloud.spanner_v1.keyset import KeyRange
 
@@ -323,7 +324,11 @@ def test_to_pb_w_only_ranges(self):
         self.assertEqual(len(result.keys), 0)
         self.assertEqual(len(result.ranges), len(RANGES))
 
-        for found, expected in zip(result.ranges, RANGES):
+        expected_ranges = [
+            KeyRangePB(start_open=KEY_1, end_closed=KEY_2),
+            KeyRangePB(start_closed=KEY_3, end_open=KEY_4),
+        ]
+        for found, expected in zip(result.ranges, expected_ranges):
             self.assertEqual(found, expected)
 
     def test_to_dict_w_all(self):
