| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| # pytype: skip-file |
| |
| import datetime |
| import logging |
| import random |
| import unittest |
| from typing import Union |
| from unittest import TestCase |
| |
| import mock |
| from bson import ObjectId |
| from bson import objectid |
| from parameterized import parameterized_class |
| from pymongo import ASCENDING |
| from pymongo import ReplaceOne |
| |
| import apache_beam as beam |
| from apache_beam.io import ReadFromMongoDB |
| from apache_beam.io import WriteToMongoDB |
| from apache_beam.io import source_test_utils |
| from apache_beam.io.mongodbio import _BoundedMongoSource |
| from apache_beam.io.mongodbio import _GenerateObjectIdFn |
| from apache_beam.io.mongodbio import _MongoSink |
| from apache_beam.io.mongodbio import _ObjectIdHelper |
| from apache_beam.io.mongodbio import _ObjectIdRangeTracker |
| from apache_beam.io.mongodbio import _WriteMongoFn |
| from apache_beam.io.range_trackers import LexicographicKeyRangeTracker |
| from apache_beam.io.range_trackers import OffsetRangeTracker |
| from apache_beam.testing.test_pipeline import TestPipeline |
| from apache_beam.testing.util import assert_that |
| from apache_beam.testing.util import equal_to |
| |
| |
| class _MockMongoColl(object): |
| """Fake mongodb collection cursor.""" |
| def __init__(self, docs): |
| self.docs = docs |
| |
| def __getitem__(self, index): |
| return self.docs[index] |
| |
| def __len__(self): |
| return len(self.docs) |
| |
| @staticmethod |
| def _make_filter(conditions): |
| assert isinstance(conditions, dict) |
| checks = [] |
| for field, value in conditions.items(): |
| if isinstance(value, dict): |
| for op, val in value.items(): |
| if op == '$gte': |
| op = '__ge__' |
| elif op == '$lt': |
| op = '__lt__' |
| else: |
| raise Exception('Operator "{0}" not supported.'.format(op)) |
| checks.append((field, op, val)) |
| else: |
| checks.append((field, '__eq__', value)) |
| |
| def func(doc): |
| for field, op, value in checks: |
| if not getattr(doc[field], op)(value): |
| return False |
| return True |
| |
| return func |
| |
| def _filter(self, filter): |
| match = [] |
| if not filter: |
| return self |
| all_filters = [] |
| if '$and' in filter: |
| for item in filter['$and']: |
| all_filters.append(self._make_filter(item)) |
| else: |
| all_filters.append(self._make_filter(filter)) |
| |
| for doc in self.docs: |
| if not all(check(doc) for check in all_filters): |
| continue |
| match.append(doc) |
| |
| return match |
| |
| @staticmethod |
| def _projection(docs, projection=None): |
| if projection: |
| return [{k: v |
| for k, v in doc.items() if k in projection or k == '_id'} |
| for doc in docs] |
| return docs |
| |
| def find(self, filter=None, projection=None, **kwargs): |
| return _MockMongoColl(self._projection(self._filter(filter), projection)) |
| |
| def sort(self, sort_items): |
| key, order = sort_items[0] |
| self.docs = sorted( |
| self.docs, key=lambda x: x[key], reverse=(order != ASCENDING)) |
| return self |
| |
| def limit(self, num): |
| return _MockMongoColl(self.docs[0:num]) |
| |
| def count_documents(self, filter): |
| return len(self._filter(filter)) |
| |
| def aggregate(self, pipeline, **kwargs): |
| # Simulate $bucketAuto aggregate pipeline. |
| # Example splits doc count for the total of 5 docs: |
| # - 1 bucket: [5] |
| # - 2 buckets: [3, 2] |
| # - 3 buckets: [2, 2, 1] |
| # - 4 buckets: [2, 1, 1, 1] |
| # - 5 buckets: [1, 1, 1, 1, 1] |
| match_step = next((step for step in pipeline if '$match' in step), None) |
| bucket_auto_step = next(step for step in pipeline if '$bucketAuto' in step) |
| if match_step is None: |
| docs = self.docs |
| else: |
| docs = self.find(filter=match_step['$match']) |
| doc_count = len(docs) |
| bucket_count = min(bucket_auto_step['$bucketAuto']['buckets'], doc_count) |
| # bucket_count ≠ 0 |
| bucket_len, remainder = divmod(doc_count, bucket_count) |
| bucket_sizes = ( |
| remainder * [bucket_len + 1] + |
| (bucket_count - remainder) * [bucket_len]) |
| buckets = [] |
| start = 0 |
| for bucket_size in bucket_sizes: |
| stop = start + bucket_size |
| if stop >= doc_count: |
| # MongoDB: the last bucket's 'max' is inclusive |
| stop = doc_count - 1 |
| count = stop - start + 1 |
| else: |
| # non-last bucket's 'max' is exclusive and == next bucket's 'min' |
| count = stop - start |
| buckets.append({ |
| '_id': { |
| 'min': docs[start]['_id'], |
| 'max': docs[stop]['_id'], |
| }, |
| 'count': count |
| }) |
| start = stop |
| |
| return buckets |
| |
| |
| class _MockMongoDb(object): |
| """Fake Mongo Db.""" |
| def __init__(self, docs): |
| self.docs = docs |
| |
| def __getitem__(self, coll_name): |
| return _MockMongoColl(self.docs) |
| |
| def command(self, command, *args, **kwargs): |
| if command == 'collstats': |
| return {'size': 5 * 1024 * 1024, 'avgObjSize': 1 * 1024 * 1024} |
| if command == 'splitVector': |
| return self.get_split_keys(command, *args, **kwargs) |
| |
| def get_split_keys(self, command, ns, min, max, maxChunkSize, **kwargs): |
| # simulate mongo db splitVector command, return split keys base on chunk |
| # size, assuming every doc is of size 1mb |
| start_id = min['_id'] |
| end_id = max['_id'] |
| if start_id >= end_id: |
| return [] |
| start_index = 0 |
| end_index = 0 |
| # get split range of [min, max] |
| for doc in self.docs: |
| if doc['_id'] < start_id: |
| start_index += 1 |
| if doc['_id'] <= end_id: |
| end_index += 1 |
| else: |
| break |
| # Return ids of elements in the range with chunk size skip and exclude |
| # head element. For simplicity of tests every document is considered 1Mb |
| # by default. |
| return { |
| 'splitKeys': [{ |
| '_id': x['_id'] |
| } for x in self.docs[start_index:end_index:maxChunkSize]][1:] |
| } |
| |
| |
| class _MockMongoClient: |
| def __init__(self, docs): |
| self.docs = docs |
| |
| def __getitem__(self, db_name): |
| return _MockMongoDb(self.docs) |
| |
| def __enter__(self): |
| return self |
| |
| def __exit__(self, exc_type, exc_val, exc_tb): |
| pass |
| |
| |
| # Generate test data for MongoDB collections of different types |
| OBJECT_IDS = [ |
| objectid.ObjectId.from_datetime( |
| datetime.datetime(year=2020, month=i + 1, day=i + 1)) for i in range(5) |
| ] |
| |
| INT_IDS = [n for n in range(5)] # [0, 1, 2, 3, 4] |
| |
| STR_IDS_1 = [str(n) for n in range(5)] # ['0', '1', '2', '3', '4'] |
| |
| # ['aaaaa', 'bbbbb', 'ccccc', 'ddddd', 'eeeee'] |
| STR_IDS_2 = [chr(97 + n) * 5 for n in range(5)] |
| |
| # ['AAAAAAAAAAAAAAAAAAAA', 'BBBBBBBBBBBBBBBBBBBB', ..., 'EEEEEEEEEEEEEEEEEEEE'] |
| STR_IDS_3 = [chr(65 + n) * 20 for n in range(5)] |
| |
| |
| @parameterized_class(('bucket_auto', '_ids', 'min_id', 'max_id'), |
| [ |
| ( |
| None, |
| OBJECT_IDS, |
| _ObjectIdHelper.int_to_id(0), |
| _ObjectIdHelper.int_to_id(2**96 - 1)), |
| ( |
| True, |
| OBJECT_IDS, |
| _ObjectIdHelper.int_to_id(0), |
| _ObjectIdHelper.int_to_id(2**96 - 1)), |
| ( |
| None, |
| INT_IDS, |
| 0, |
| 2**96 - 1, |
| ), |
| ( |
| True, |
| INT_IDS, |
| 0, |
| 2**96 - 1, |
| ), |
| ( |
| None, |
| STR_IDS_1, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ( |
| True, |
| STR_IDS_1, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ( |
| None, |
| STR_IDS_2, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ( |
| True, |
| STR_IDS_2, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ( |
| None, |
| STR_IDS_3, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ( |
| True, |
| STR_IDS_3, |
| chr(0), |
| chr(0x10ffff), |
| ), |
| ]) |
| class MongoSourceTest(unittest.TestCase): |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def setUp(self, mock_client): |
| self._docs = [{'_id': self._ids[i], 'x': i} for i in range(len(self._ids))] |
| mock_client.return_value = _MockMongoClient(self._docs) |
| |
| self.mongo_source = self._create_source(bucket_auto=self.bucket_auto) |
| |
| @staticmethod |
| def _create_source(filter=None, bucket_auto=None): |
| kwargs = {} |
| if filter is not None: |
| kwargs['filter'] = filter |
| if bucket_auto is not None: |
| kwargs['bucket_auto'] = bucket_auto |
| return _BoundedMongoSource('mongodb://test', 'testdb', 'testcoll', **kwargs) |
| |
| def _increment_id( |
| self, |
| _id: Union[ObjectId, int, str], |
| inc: int, |
| ) -> Union[ObjectId, int, str]: |
| """Helper method to increment `_id` of different types.""" |
| |
| if isinstance(_id, ObjectId): |
| return _ObjectIdHelper.increment_id(_id, inc) |
| |
| if isinstance(_id, int): |
| return _id + inc |
| |
| if isinstance(_id, str): |
| index = self._ids.index(_id) + inc |
| if index <= 0: |
| return self._ids[0] |
| if index >= len(self._ids): |
| return self._ids[-1] |
| return self._ids[index] |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_estimate_size(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs) |
| self.assertEqual(self.mongo_source.estimate_size(), 5 * 1024 * 1024) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_estimate_average_document_size(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs) |
| self.assertEqual( |
| self.mongo_source._estimate_average_document_size(), 1 * 1024 * 1024) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_split(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs) |
| for size_mb, expected_split_count in [(0.5, 5), (1, 5), (2, 3), (10, 1)]: |
| size = size_mb * 1024 * 1024 |
| splits = list( |
| self.mongo_source.split( |
| start_position=None, stop_position=None, |
| desired_bundle_size=size)) |
| |
| self.assertEqual(len(splits), expected_split_count) |
| reference_info = (self.mongo_source, None, None) |
| sources_info = ([ |
| (split.source, split.start_position, split.stop_position) |
| for split in splits |
| ]) |
| source_test_utils.assert_sources_equal_reference_source( |
| reference_info, sources_info) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_split_single_document(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs[0:1]) |
| for size_mb in [1, 5]: |
| size = size_mb * 1024 * 1024 |
| splits = list( |
| self.mongo_source.split( |
| start_position=None, stop_position=None, |
| desired_bundle_size=size)) |
| self.assertEqual(len(splits), 1) |
| _id = self._docs[0]['_id'] |
| assert _id == splits[0].start_position |
| assert _id <= splits[0].stop_position |
| if isinstance(_id, (ObjectId, int)): |
| # We can unambiguously determine next `_id` |
| assert self._increment_id(_id, 1) == splits[0].stop_position |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_split_no_documents(self, mock_client): |
| mock_client.return_value = _MockMongoClient([]) |
| with self.assertRaises(ValueError) as cm: |
| list( |
| self.mongo_source.split( |
| start_position=None, |
| stop_position=None, |
| desired_bundle_size=1024 * 1024)) |
| self.assertEqual(str(cm.exception), 'Empty Mongodb collection') |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_split_filtered(self, mock_client): |
| # filtering 2 documents: 2 <= 'x' < 4 |
| filtered_mongo_source = self._create_source( |
| filter={'x': { |
| '$gte': 2, '$lt': 4 |
| }}, bucket_auto=self.bucket_auto) |
| |
| mock_client.return_value = _MockMongoClient(self._docs) |
| for size_mb, (bucket_auto_count, split_vector_count) in [(1, (2, 5)), |
| (2, (1, 3)), |
| (10, (1, 1))]: |
| size = size_mb * 1024 * 1024 |
| splits = list( |
| filtered_mongo_source.split( |
| start_position=None, stop_position=None, |
| desired_bundle_size=size)) |
| |
| if self.bucket_auto: |
| self.assertEqual(len(splits), bucket_auto_count) |
| else: |
| # Note: splitVector mode does not respect filter |
| self.assertEqual(len(splits), split_vector_count) |
| reference_info = ( |
| filtered_mongo_source, self._docs[2]['_id'], self._docs[4]['_id']) |
| sources_info = ([ |
| (split.source, split.start_position, split.stop_position) |
| for split in splits |
| ]) |
| source_test_utils.assert_sources_equal_reference_source( |
| reference_info, sources_info) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_split_filtered_empty(self, mock_client): |
| # filtering doesn't match any documents |
| filtered_mongo_source = self._create_source( |
| filter={'x': { |
| '$lt': 0 |
| }}, bucket_auto=self.bucket_auto) |
| |
| mock_client.return_value = _MockMongoClient(self._docs) |
| for size_mb, (bucket_auto_count, split_vector_count) in [(1, (1, 5)), |
| (2, (1, 3)), |
| (10, (1, 1))]: |
| size = size_mb * 1024 * 1024 |
| splits = list( |
| filtered_mongo_source.split( |
| start_position=None, stop_position=None, |
| desired_bundle_size=size)) |
| |
| if self.bucket_auto: |
| # Note: if filter matches no docs - one split covers entire range |
| self.assertEqual(len(splits), bucket_auto_count) |
| else: |
| # Note: splitVector mode does not respect filter |
| self.assertEqual(len(splits), split_vector_count) |
| reference_info = ( |
| filtered_mongo_source, |
| # range to match no documents: |
| self._increment_id(self._docs[-1]['_id'], 1), |
| self._increment_id(self._docs[-1]['_id'], 2), |
| ) |
| sources_info = ([ |
| (split.source, split.start_position, split.stop_position) |
| for split in splits |
| ]) |
| source_test_utils.assert_sources_equal_reference_source( |
| reference_info, sources_info) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_dynamic_work_rebalancing(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs) |
| splits = list( |
| self.mongo_source.split(desired_bundle_size=3000 * 1024 * 1024)) |
| assert len(splits) == 1 |
| source_test_utils.assert_split_at_fraction_exhaustive( |
| splits[0].source, splits[0].start_position, splits[0].stop_position) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_get_range_tracker(self, mock_client): |
| mock_client.return_value = _MockMongoClient(self._docs) |
| if self._ids == OBJECT_IDS: |
| self.assertIsInstance( |
| self.mongo_source.get_range_tracker(None, None), |
| _ObjectIdRangeTracker, |
| ) |
| elif self._ids == INT_IDS: |
| self.assertIsInstance( |
| self.mongo_source.get_range_tracker(None, None), |
| OffsetRangeTracker, |
| ) |
| elif self._ids == STR_IDS_1: |
| self.assertIsInstance( |
| self.mongo_source.get_range_tracker(None, None), |
| LexicographicKeyRangeTracker, |
| ) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_read(self, mock_client): |
| mock_tracker = mock.MagicMock() |
| test_cases = [ |
| { |
| # range covers the first(inclusive) to third(exclusive) documents |
| 'start': self._ids[0], |
| 'stop': self._ids[2], |
| 'expected': self._docs[0:2] |
| }, |
| { |
| # range covers from the first to the third documents |
| 'start': self.min_id, # smallest possible id |
| 'stop': self._ids[2], |
| 'expected': self._docs[0:2] |
| }, |
| { |
| # range covers from the third to last documents |
| 'start': self._ids[2], |
| 'stop': self.max_id, # largest possible id |
| 'expected': self._docs[2:] |
| }, |
| { |
| # range covers all documents |
| 'start': self.min_id, |
| 'stop': self.max_id, |
| 'expected': self._docs |
| }, |
| { |
| # range doesn't include any document |
| 'start': self._increment_id(self._ids[2], 1), |
| 'stop': self._increment_id(self._ids[3], -1), |
| 'expected': [] |
| }, |
| ] |
| mock_client.return_value = _MockMongoClient(self._docs) |
| for case in test_cases: |
| mock_tracker.start_position.return_value = case['start'] |
| mock_tracker.stop_position.return_value = case['stop'] |
| result = list(self.mongo_source.read(mock_tracker)) |
| self.assertListEqual(case['expected'], result) |
| |
| def test_display_data(self): |
| data = self.mongo_source.display_data() |
| self.assertTrue('database' in data) |
| self.assertTrue('collection' in data) |
| |
| def test_range_is_not_splittable(self): |
| self.assertTrue( |
| self.mongo_source._range_is_not_splittable( |
| _ObjectIdHelper.int_to_id(1), |
| _ObjectIdHelper.int_to_id(1), |
| )) |
| self.assertTrue( |
| self.mongo_source._range_is_not_splittable( |
| _ObjectIdHelper.int_to_id(1), |
| _ObjectIdHelper.int_to_id(2), |
| )) |
| self.assertFalse( |
| self.mongo_source._range_is_not_splittable( |
| _ObjectIdHelper.int_to_id(1), |
| _ObjectIdHelper.int_to_id(3), |
| )) |
| |
| self.assertTrue(self.mongo_source._range_is_not_splittable(0, 0)) |
| self.assertTrue(self.mongo_source._range_is_not_splittable(0, 1)) |
| self.assertFalse(self.mongo_source._range_is_not_splittable(0, 2)) |
| |
| self.assertTrue(self.mongo_source._range_is_not_splittable("AAA", "AAA")) |
| self.assertFalse( |
| self.mongo_source._range_is_not_splittable("AAA", "AAA\x00")) |
| self.assertFalse(self.mongo_source._range_is_not_splittable("AAA", "AAB")) |
| |
| |
| @parameterized_class(('bucket_auto', ), [(False, ), (True, )]) |
| class ReadFromMongoDBTest(unittest.TestCase): |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_read_from_mongodb(self, mock_client): |
| documents = [{ |
| '_id': objectid.ObjectId(), 'x': i, 'selected': 1, 'unselected': 2 |
| } for i in range(3)] |
| mock_client.return_value = _MockMongoClient(documents) |
| |
| projection = ['x', 'selected'] |
| projected_documents = [{ |
| k: v |
| for k, v in e.items() if k in projection or k == '_id' |
| } for e in documents] |
| |
| with TestPipeline() as p: |
| docs = p | 'ReadFromMongoDB' >> ReadFromMongoDB( |
| uri='mongodb://test', |
| db='db', |
| coll='collection', |
| projection=projection, |
| bucket_auto=self.bucket_auto) |
| assert_that(docs, equal_to(projected_documents)) |
| |
| |
| class GenerateObjectIdFnTest(unittest.TestCase): |
| def test_process(self): |
| with TestPipeline() as p: |
| output = ( |
| p | "Create" >> beam.Create([{ |
| 'x': 1 |
| }, { |
| 'x': 2, '_id': 123 |
| }]) |
| | "Generate ID" >> beam.ParDo(_GenerateObjectIdFn()) |
| | "Check" >> beam.Map(lambda x: '_id' in x)) |
| assert_that(output, equal_to([True] * 2)) |
| |
| |
| class WriteMongoFnTest(unittest.TestCase): |
| @mock.patch('apache_beam.io.mongodbio._MongoSink') |
| def test_process(self, mock_sink): |
| docs = [{'x': 1}, {'x': 2}, {'x': 3}] |
| with TestPipeline() as p: |
| _ = ( |
| p | "Create" >> beam.Create(docs) |
| | "Write" >> beam.ParDo(_WriteMongoFn(batch_size=2))) |
| p.run() |
| |
| self.assertEqual( |
| 2, mock_sink.return_value.__enter__.return_value.write.call_count) |
| |
| def test_display_data(self): |
| data = _WriteMongoFn(batch_size=10).display_data() |
| self.assertEqual(10, data['batch_size']) |
| |
| |
| class MongoSinkTest(unittest.TestCase): |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_write(self, mock_client): |
| docs = [{'x': 1}, {'x': 2}, {'x': 3}] |
| _MongoSink(uri='test', db='test', coll='test').write(docs) |
| self.assertTrue( |
| mock_client.return_value.__getitem__.return_value.__getitem__. |
| return_value.bulk_write.called) |
| |
| |
| class WriteToMongoDBTest(unittest.TestCase): |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_write_to_mongodb_with_existing_id(self, mock_client): |
| _id = objectid.ObjectId() |
| docs = [{'x': 1, '_id': _id}] |
| expected_update = [ |
| ReplaceOne({'_id': _id}, { |
| 'x': 1, '_id': _id |
| }, True, None) |
| ] |
| with TestPipeline() as p: |
| _ = ( |
| p | "Create" >> beam.Create(docs) |
| | "Write" >> WriteToMongoDB(db='test', coll='test')) |
| p.run() |
| mock_client.return_value.__getitem__.return_value.__getitem__. \ |
| return_value.bulk_write.assert_called_with(expected_update) |
| |
| @mock.patch('apache_beam.io.mongodbio.MongoClient') |
| def test_write_to_mongodb_with_generated_id(self, mock_client): |
| docs = [{'x': 1}] |
| expected_update = [ |
| ReplaceOne({'_id': mock.ANY}, { |
| 'x': 1, '_id': mock.ANY |
| }, True, None) |
| ] |
| with TestPipeline() as p: |
| _ = ( |
| p | "Create" >> beam.Create(docs) |
| | "Write" >> WriteToMongoDB(db='test', coll='test')) |
| p.run() |
| mock_client.return_value.__getitem__.return_value.__getitem__. \ |
| return_value.bulk_write.assert_called_with(expected_update) |
| |
| |
| class ObjectIdHelperTest(TestCase): |
| def test_conversion(self): |
| test_cases = [ |
| (objectid.ObjectId('000000000000000000000000'), 0), |
| (objectid.ObjectId('000000000000000100000000'), 2**32), |
| (objectid.ObjectId('0000000000000000ffffffff'), 2**32 - 1), |
| (objectid.ObjectId('000000010000000000000000'), 2**64), |
| (objectid.ObjectId('00000000ffffffffffffffff'), 2**64 - 1), |
| (objectid.ObjectId('ffffffffffffffffffffffff'), 2**96 - 1), |
| ] |
| for (_id, number) in test_cases: |
| self.assertEqual(_id, _ObjectIdHelper.int_to_id(number)) |
| self.assertEqual(number, _ObjectIdHelper.id_to_int(_id)) |
| |
| # random tests |
| for _ in range(100): |
| _id = objectid.ObjectId() |
| number = int(_id.binary.hex(), 16) |
| self.assertEqual(_id, _ObjectIdHelper.int_to_id(number)) |
| self.assertEqual(number, _ObjectIdHelper.id_to_int(_id)) |
| |
| def test_increment_id(self): |
| test_cases = [ |
| ( |
| objectid.ObjectId("000000000000000100000000"), |
| objectid.ObjectId("0000000000000000ffffffff"), |
| ), |
| ( |
| objectid.ObjectId("000000010000000000000000"), |
| objectid.ObjectId("00000000ffffffffffffffff"), |
| ), |
| ] |
| for first, second in test_cases: |
| self.assertEqual(second, _ObjectIdHelper.increment_id(first, -1)) |
| self.assertEqual(first, _ObjectIdHelper.increment_id(second, 1)) |
| |
| for _ in range(100): |
| _id = objectid.ObjectId() |
| self.assertLess(_id, _ObjectIdHelper.increment_id(_id, 1)) |
| self.assertGreater(_id, _ObjectIdHelper.increment_id(_id, -1)) |
| |
| |
| class ObjectRangeTrackerTest(TestCase): |
| def test_fraction_position_conversion(self): |
| start_int = 0 |
| stop_int = 2**96 - 1 |
| start = _ObjectIdHelper.int_to_id(start_int) |
| stop = _ObjectIdHelper.int_to_id(stop_int) |
| test_cases = ([start_int, stop_int, 2**32, 2**32 - 1, 2**64, 2**64 - 1] + |
| [random.randint(start_int, stop_int) for _ in range(100)]) |
| tracker = _ObjectIdRangeTracker() |
| for pos in test_cases: |
| _id = _ObjectIdHelper.int_to_id(pos - start_int) |
| desired_fraction = (pos - start_int) / (stop_int - start_int) |
| self.assertAlmostEqual( |
| tracker.position_to_fraction(_id, start, stop), |
| desired_fraction, |
| places=20) |
| |
| convert_id = tracker.fraction_to_position( |
| (pos - start_int) / (stop_int - start_int), start, stop) |
| # due to precision loss, the convert fraction is only gonna be close to |
| # original fraction. |
| convert_fraction = tracker.position_to_fraction(convert_id, start, stop) |
| |
| self.assertGreater(convert_id, start) |
| self.assertLess(convert_id, stop) |
| self.assertAlmostEqual(convert_fraction, desired_fraction, places=20) |
| |
| |
| if __name__ == '__main__': |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
