import pymongo import pytest from pymongo.errors import DocumentTooLarge, OperationFailure from maggma.stores import GridFSStore, MemoryStore, MongoStore from maggma.stores.shared_stores import MultiStore, StoreFacade from maggma.validators import JSONSchemaValidator @pytest.fixture() def mongostore(): store = MongoStore("maggma_test", "test") store.connect() yield store store._collection.drop() @pytest.fixture() def gridfsstore(): store = GridFSStore("maggma_test", "test", key="task_id") store.connect() yield store store._files_collection.drop() store._chunks_collection.drop() @pytest.fixture() def multistore(): return MultiStore() @pytest.fixture() def memorystore(): store = MemoryStore() store.connect() return store def test_add_stores(multistore, mongostore, gridfsstore): # Should be empty at the start assert multistore.count_stores() == 0 multistore.ensure_store(mongostore) assert multistore.count_stores() == 1 assert multistore.get_store_index(mongostore) == 0 # Attempting to reinsert this store should do nothing multistore.ensure_store(mongostore) assert multistore.count_stores() == 1 # Make a copy of the mongostore and it should still do nothing temp_mongostore = MongoStore.from_dict(mongostore.as_dict()) multistore.ensure_store(temp_mongostore) assert multistore.count_stores() == 1 assert multistore.get_store_index(temp_mongostore) == 0 # Add this copy again, but don't use ensure_store # This tests the case in which a prior thread added # the store, but this current process was already # waiting for the lock acquisition multistore.add_store(temp_mongostore) assert multistore.count_stores() == 1 # Add the GridFSStore to the MultiStore() multistore.ensure_store(gridfsstore) assert multistore.count_stores() == 2 assert multistore.get_store_index(gridfsstore) == 1 # Add something that isn't a store class DummyObject: def __init__(self, a: int): self.a = a with pytest.raises(TypeError): multistore.ensure_store(DummyObject(1)) def test_store_facade(multistore, mongostore, gridfsstore): StoreFacade(mongostore, multistore) assert multistore.count_stores() == 1 assert multistore.get_store_index(mongostore) == 0 StoreFacade(gridfsstore, multistore) assert multistore.count_stores() == 2 assert multistore.get_store_index(gridfsstore) == 1 def test_multistore_query(multistore, mongostore, memorystore): memorystore_facade = StoreFacade(memorystore, multistore) mongostore_facade = StoreFacade(mongostore, multistore) temp_mongostore_facade = StoreFacade(MongoStore.from_dict(mongostore.as_dict()), multistore) memorystore_facade._collection.insert_one({"a": 1, "b": 2, "c": 3}) assert memorystore_facade.query_one(properties=["a"])["a"] == 1 assert memorystore_facade.query_one(properties=["a"])["a"] == 1 assert memorystore_facade.query_one(properties=["b"])["b"] == 2 assert memorystore_facade.query_one(properties=["c"])["c"] == 3 mongostore_facade._collection.insert_one({"a": 4, "b": 5, "c": 6}) assert mongostore_facade.query_one(properties=["a"])["a"] == 4 assert mongostore_facade.query_one(properties=["a"])["a"] == 4 assert mongostore_facade.query_one(properties=["b"])["b"] == 5 assert mongostore_facade.query_one(properties=["c"])["c"] == 6 assert temp_mongostore_facade.query_one(properties=["a"])["a"] == 4 assert temp_mongostore_facade.query_one(properties=["a"])["a"] == 4 assert temp_mongostore_facade.query_one(properties=["b"])["b"] == 5 assert temp_mongostore_facade.query_one(properties=["c"])["c"] == 6 def test_multistore_count(multistore, mongostore, memorystore): memorystore_facade = StoreFacade(memorystore, multistore) memorystore_facade._collection.insert_one({"a": 1, "b": 2, "c": 3}) assert memorystore_facade.count() == 1 memorystore_facade._collection.insert_one({"aa": 1, "b": 2, "c": 3}) assert memorystore_facade.count() == 2 assert memorystore_facade.count({"a": 1}) == 1 def test_multistore_distinct(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) mongostore_facade._collection.insert_one({"a": 1, "b": 2, "c": 3}) mongostore_facade._collection.insert_one({"a": 4, "d": 5, "e": 6, "g": {"h": 1}}) assert set(mongostore_facade.distinct("a")) == {1, 4} # Test list distinct functionality mongostore_facade._collection.insert_one({"a": 4, "d": 6, "e": 7}) mongostore_facade._collection.insert_one({"a": 4, "d": 6, "g": {"h": 2}}) # Test distinct subdocument functionality ghs = mongostore_facade.distinct("g.h") assert set(ghs) == {1, 2} # Test when key doesn't exist assert mongostore_facade.distinct("blue") == [] # Test when null is a value mongostore_facade._collection.insert_one({"i": None}) assert mongostore_facade.distinct("i") == [None] # Test to make sure DocumentTooLarge errors get dealt with properly using built in distinct mongostore_facade._collection.insert_many([{"key": [f"mp-{i}"]} for i in range(1000000)]) vals = mongostore_facade.distinct("key") # Test to make sure distinct on array field is unraveled when using manual distinct assert len(vals) == len(list(range(1000000))) assert all(isinstance(v, str) for v in vals) # Test to make sure manual distinct uses the criteria query mongostore_facade._collection.insert_many([{"key": f"mp-{i}", "a": 2} for i in range(1000001, 2000001)]) vals = mongostore_facade.distinct("key", {"a": 2}) assert len(vals) == len(list(range(1000001, 2000001))) def test_multistore_update(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) mongostore_facade.update({"e": 6, "d": 4}, key="e") assert mongostore_facade.query_one(criteria={"d": {"$exists": 1}}, properties=["d"])["d"] == 4 mongostore_facade.update([{"e": 7, "d": 8, "f": 9}], key=["d", "f"]) assert mongostore_facade.query_one(criteria={"d": 8, "f": 9}, properties=["e"])["e"] == 7 mongostore_facade.update([{"e": 11, "d": 8, "f": 9}], key=["d", "f"]) assert mongostore_facade.query_one(criteria={"d": 8, "f": 9}, properties=["e"])["e"] == 11 test_schema = { "type": "object", "properties": {"e": {"type": "integer"}}, "required": ["e"], } mongostore_facade.validator = JSONSchemaValidator(schema=test_schema) mongostore_facade.update({"e": 100, "d": 3}, key="e") # Continue to update doc when validator is not set to strict mode mongostore_facade.update({"e": "abc", "d": 3}, key="e") # ensure safe_update works to not throw DocumentTooLarge errors large_doc = {f"mp-{i}": f"mp-{i}" for i in range(1000000)} large_doc["e"] = 999 with pytest.raises((OperationFailure, DocumentTooLarge)): mongostore_facade.update([large_doc, {"e": 1001}], key="e") mongostore_facade.safe_update = True assert mongostore_facade.safe_update is True mongostore_facade.update([large_doc, {"e": 1001}], key="e") assert mongostore_facade.query_one({"e": 1001}) is not None def test_multistore_groupby(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) mongostore_facade.update( [ {"e": 7, "d": 9, "f": 9}, {"e": 7, "d": 9, "f": 10}, {"e": 8, "d": 9, "f": 11}, {"e": 9, "d": 10, "f": 12}, ], key="f", ) data = list(mongostore_facade.groupby("d")) assert len(data) == 2 grouped_by_9 = next(g[1] for g in data if g[0]["d"] == 9) assert len(grouped_by_9) == 3 grouped_by_10 = next(g[1] for g in data if g[0]["d"] == 10) assert len(grouped_by_10) == 1 data = list(mongostore_facade.groupby(["e", "d"])) assert len(data) == 3 def test_multistore_remove_docs(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) mongostore_facade._collection.insert_one({"a": 1, "b": 2, "c": 3}) mongostore_facade._collection.insert_one({"a": 4, "d": 5, "e": 6, "g": {"h": 1}}) mongostore_facade.remove_docs({"a": 1}) assert len(list(mongostore_facade.query({"a": 4}))) == 1 assert len(list(mongostore_facade.query({"a": 1}))) == 0 def test_multistore_connect_reconnect(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) assert isinstance(mongostore_facade._collection, pymongo.collection.Collection) mongostore_facade.close() assert mongostore_facade._coll is None mongostore_facade.connect() # Test using the multistore to close connections multistore.close_all() assert mongostore_facade._coll is None multistore.connect_all() assert isinstance(mongostore_facade._collection, pymongo.collection.Collection) def test_multistore_name(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) assert mongostore_facade.name == "mongo://localhost/maggma_test/test" def test_multistore_ensure_index(multistore, mongostore): mongostore_facade = StoreFacade(mongostore, multistore) assert mongostore_facade.ensure_index("test_key") # TODO: How to check for exception?