1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
"""
Tests for Projection_Builder
"""
import pytest
from maggma.builders.projection_builder import Projection_Builder
from maggma.stores import MemoryStore
@pytest.fixture()
def source1():
store = MemoryStore("source1", key="k", last_updated_field="lu")
store.connect()
store.ensure_index("k")
store.ensure_index("lu")
store.update([{"k": k, "a": "a", "b": "b"} for k in range(10)])
return store
@pytest.fixture()
def source2():
store = MemoryStore("source2", key="k", last_updated_field="lu")
store.connect()
store.ensure_index("k")
store.ensure_index("lu")
store.update([{"k": k, "c": "c", "d": "d"} for k in range(15)])
return store
@pytest.fixture()
def target():
store = MemoryStore("target", key="k", last_updated_field="lu")
store.connect()
store.ensure_index("k")
store.ensure_index("lu")
return store
def test_get_items(source1, source2, target):
builder = Projection_Builder(source_stores=[source1, source2], target_store=target)
items = next(iter(builder.get_items()))
assert len(items) == 25
def test_process_item(source1, source2, target):
# test fields_to_project = empty dict and list
builder = Projection_Builder(
source_stores=[source1, source2],
target_store=target,
fields_to_project=[[], {}],
)
items = next(iter(builder.get_items()))
outputs = builder.process_item(items)
assert len(outputs) == 15
output = next(o for o in outputs if o["k"] < 10)
assert all(k in ["k", "a", "b", "c", "d"] for k in output)
output = next(o for o in outputs if o["k"] > 9)
assert all(k in ["k", "c", "d"] for k in output)
assert all(k not in ["a", "b"] for k in output)
# test fields_to_project = lists
builder = Projection_Builder(
source_stores=[source1, source2],
target_store=target,
fields_to_project=[["a", "b"], ["d"]],
)
items = next(iter(builder.get_items()))
outputs = builder.process_item(items)
output = next(o for o in outputs if o["k"] < 10)
assert all(k in ["k", "a", "b", "d"] for k in output)
assert all(k not in ["c"] for k in output)
# test fields_to_project = dict and list
builder = Projection_Builder(
source_stores=[source1, source2],
target_store=target,
fields_to_project=[{"newa": "a", "b": "b"}, ["d"]],
)
items = next(iter(builder.get_items()))
outputs = builder.process_item(items)
output = next(o for o in outputs if o["k"] < 10)
assert all(k in ["k", "newa", "b", "d"] for k in output)
assert all(k not in ["a", "c"] for k in output)
def test_update_targets(source1, source2, target):
builder = Projection_Builder(
source_stores=[source1, source2],
target_store=target,
fields_to_project=[{"newa": "a", "b": "b"}, ["d"]],
)
items = list(builder.get_items())
processed_chunk = [builder.process_item(item) for item in items]
processed_items = [item for item in processed_chunk if item is not None]
builder.update_targets(processed_items)
assert len(list(target.query())) == 15
assert target.query_one(criteria={"k": 0})["newa"] == "a"
assert target.query_one(criteria={"k": 0})["d"] == "d"
assert target.query_one(criteria={"k": 10})["d"] == "d"
assert "a" not in target.query_one(criteria={"k": 10})
def test_run(source1, source2, target):
builder = Projection_Builder(source_stores=[source1, source2], target_store=target)
builder.run()
assert len(list(target.query())) == 15
assert target.query_one(criteria={"k": 0})["a"] == "a"
assert target.query_one(criteria={"k": 0})["d"] == "d"
assert target.query_one(criteria={"k": 10})["d"] == "d"
assert "a" not in target.query_one(criteria={"k": 10})
def test_query(source1, source2, target):
target.remove_docs({})
builder = Projection_Builder(
source_stores=[source1, source2],
target_store=target,
query_by_key=[0, 1, 2, 3, 4],
)
builder.run()
assert len(list(target.query())) == 5
|
