Repository: beam Updated Branches: refs/heads/master 754d5d833 -> aad85553b
Use built-in cmp python function in comparing datastore paths Project: http://git-wip-us.apache.org/repos/asf/beam/repo Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/0a17c3da Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/0a17c3da Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/0a17c3da Branch: refs/heads/master Commit: 0a17c3da5a249dbe4408c5dce7e1825897865e1e Parents: 754d5d8 Author: Vikas Kedigehalli <vika...@google.com> Authored: Mon May 15 15:50:25 2017 -0700 Committer: Ahmet Altay <al...@google.com> Committed: Mon May 15 18:42:00 2017 -0700 ---------------------------------------------------------------------- .../io/gcp/datastore/v1/fake_datastore.py | 8 +++-- .../apache_beam/io/gcp/datastore/v1/helper.py | 14 ++------ .../io/gcp/datastore/v1/query_splitter_test.py | 37 +++++++++++--------- 3 files changed, 30 insertions(+), 29 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/beam/blob/0a17c3da/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py b/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py index 2332579..aa37805 100644 --- a/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py +++ b/sdks/python/apache_beam/io/gcp/datastore/v1/fake_datastore.py @@ -90,13 +90,17 @@ def create_response(entities, end_cursor, finish): return resp -def create_entities(count): +def create_entities(count, id_or_name=False): """Creates a list of entities with random keys.""" entities = [] for _ in range(count): entity_result = query_pb2.EntityResult() - entity_result.entity.key.path.add().name = str(uuid.uuid4()) + if id_or_name: + entity_result.entity.key.path.add().id = ( + uuid.uuid4().int & ((1 << 63) - 1)) + else: + entity_result.entity.key.path.add().name = str(uuid.uuid4()) entities.append(entity_result) return entities http://git-wip-us.apache.org/repos/asf/beam/blob/0a17c3da/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py b/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py index 9e2c053..f977536 100644 --- a/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py +++ b/sdks/python/apache_beam/io/gcp/datastore/v1/helper.py @@ -80,7 +80,7 @@ def compare_path(p1, p2): 3. If no `id` is defined for both paths, then their `names` are compared. """ - result = str_compare(p1.kind, p2.kind) + result = cmp(p1.kind, p2.kind) if result != 0: return result @@ -88,20 +88,12 @@ def compare_path(p1, p2): if not p2.HasField('id'): return -1 - return p1.id - p2.id + return cmp(p1.id, p2.id) if p2.HasField('id'): return 1 - return str_compare(p1.name, p2.name) - - -def str_compare(s1, s2): - if s1 == s2: - return 0 - elif s1 < s2: - return -1 - return 1 + return cmp(p1.name, p2.name) def get_datastore(project): http://git-wip-us.apache.org/repos/asf/beam/blob/0a17c3da/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py ---------------------------------------------------------------------- diff --git a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py index b7b054f..52f25fa 100644 --- a/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py +++ b/sdks/python/apache_beam/io/gcp/datastore/v1/query_splitter_test.py @@ -157,27 +157,32 @@ class QuerySplitterTest(unittest.TestCase): batch_size: the number of entities returned by fake datastore in one req. """ - entities = fake_datastore.create_entities(num_entities) - mock_datastore = MagicMock() - # Assign a fake run_query method as a side_effect to the mock. - mock_datastore.run_query.side_effect = \ - fake_datastore.create_run_query(entities, batch_size) + # Test for both random long ids and string ids. + id_or_name = [True, False] - split_queries = query_splitter.get_splits(mock_datastore, query, num_splits) + for id_type in id_or_name: + entities = fake_datastore.create_entities(num_entities, id_type) + mock_datastore = MagicMock() + # Assign a fake run_query method as a side_effect to the mock. + mock_datastore.run_query.side_effect = \ + fake_datastore.create_run_query(entities, batch_size) - # if request num_splits is greater than num_entities, the best it can - # do is one entity per split. - expected_num_splits = min(num_splits, num_entities + 1) - self.assertEqual(len(split_queries), expected_num_splits) + split_queries = query_splitter.get_splits( + mock_datastore, query, num_splits) - expected_requests = QuerySplitterTest.create_scatter_requests( - query, num_splits, batch_size, num_entities) + # if request num_splits is greater than num_entities, the best it can + # do is one entity per split. + expected_num_splits = min(num_splits, num_entities + 1) + self.assertEqual(len(split_queries), expected_num_splits) - expected_calls = [] - for req in expected_requests: - expected_calls.append(call(req)) + expected_requests = QuerySplitterTest.create_scatter_requests( + query, num_splits, batch_size, num_entities) - self.assertEqual(expected_calls, mock_datastore.run_query.call_args_list) + expected_calls = [] + for req in expected_requests: + expected_calls.append(call(req)) + + self.assertEqual(expected_calls, mock_datastore.run_query.call_args_list) @staticmethod def create_scatter_requests(query, num_splits, batch_size, num_entities):
