def test_source_on_clone():
assert {
"_source": {
"includes": ["foo.bar.*"],
"excludes": ["foo.one"]
},
"query": {
"bool": {
"filter": [{
"term": {
"title": {
"value": "python"
}
}
}]
}
},
} == Search().source(includes=["foo.bar.*"]).source(
excludes=["foo.one"]).filter("term", title="python").to_dict()
assert {
"_source": False,
"query": {
"bool": {
"filter": [{
"term": {
"title": {
"value": "python"
}
}
}]
}
},
} == Search().source(False).filter("term", title="python").to_dict()
def test_search_index(self):
s = Search(index="i")
assert s._index == ["i"]
s = s.index("i2")
assert s._index == ["i", "i2"]
s = s.index(u"i3")
assert s._index == ["i", "i2", "i3"]
s = s.index()
assert s._index is None
s = Search(index=("i", "i2"))
assert s._index == ["i", "i2"]
s = Search(index=["i", "i2"])
assert s._index == ["i", "i2"]
s = Search()
s = s.index("i", "i2")
assert s._index == ["i", "i2"]
s2 = s.index("i3")
assert s._index == ["i", "i2"]
assert s2._index == ["i", "i2", "i3"]
s = Search()
s = s.index(["i", "i2"], "i3")
assert s._index == ["i", "i2", "i3"]
s2 = s.index("i4")
assert s._index == ["i", "i2", "i3"]
assert s2._index == ["i", "i2", "i3", "i4"]
s2 = s.index(["i4"])
assert s2._index == ["i", "i2", "i3", "i4"]
s2 = s.index(("i4", "i5"))
assert s2._index == ["i", "i2", "i3", "i4", "i5"]
def test_sort():
s = Search()
s = s.sort("fielda", "-fieldb")
assert ["fielda", {"fieldb": {"order": "desc"}}] == s._sort
assert {"sort": ["fielda", {"fieldb": {"order": "desc"}}]} == s.to_dict()
s = s.sort()
assert [] == s._sort
assert Search().to_dict() == s.to_dict()
def test_source_copied_on_clone(self):
s = Search().source(False)
self.assertEqual(s._clone()._source, s._source)
self.assertIs(s._clone()._source, False)
s2 = Search().source([])
assert s2._clone()._source == s2._source
assert s2._source == []
s3 = Search().source(["some", "fields"])
assert s3._clone()._source == s3._source
assert s3._clone()._source == ["some", "fields"]
def test_source_copied_on_clone():
s = Search().source(False)
assert s._clone()._source == s._source
assert s._clone()._source is False
s2 = Search().source([])
assert s2._clone()._source == s2._source
assert s2._source == []
s3 = Search().source(["some", "fields"])
assert s3._clone()._source == s3._source
assert s3._clone()._source == ["some", "fields"]
def test_search_to_dict(self):
s = Search()
assert {} == s.to_dict()
s = s.query("match", f=42)
assert {"query": {"match": {"f": {"query": 42}}}} == s.to_dict()
assert {
"query": {
"match": {
"f": {
"query": 42
}
}
},
"size": 10
} == s.to_dict(size=10)
s = s.aggs("per_tag", "terms", field="f").aggs("max_score",
"max",
field="score")
d = {
"aggs": {
"per_tag": {
"terms": {
"field": "f"
},
"aggs": {
"max_score": {
"max": {
"field": "score"
}
}
},
}
},
"query": {
"match": {
"f": {
"query": 42
}
}
},
}
self.assertEqual(d, s.to_dict())
s = Search().params(size=5)
assert {"size": 5} == s.to_dict()
s = s.params(from_=42)
assert {"size": 5, "from": 42} == s.to_dict()
def search(self, nested_autocorrect=True):
return Search(
using=self.client,
mapping=self._mapping,
index=self.name,
nested_autocorrect=nested_autocorrect,
)
def test_using(self):
o = object()
o2 = object()
s = Search(using=o)
self.assertIs(s._using, o)
s2 = s.using(o2)
self.assertIs(s._using, o)
self.assertIs(s2._using, o2)
def test_copy_clones(self):
from copy import copy
s1 = Search().source(["some", "fields"])
s2 = copy(s1)
assert s1 == s2
assert s1 is not s2
def test_using():
o = object()
o2 = object()
s = Search(using=o)
assert s._using is o
s2 = s.using(o2)
assert s._using is o
assert s2._using is o2
def test_source():
assert {} == Search().source().to_dict()
assert {
"_source": {
"includes": ["foo.bar.*"],
"excludes": ["foo.one"]
}
} == Search().source(includes=["foo.bar.*"],
excludes=["foo.one"]).to_dict()
assert {"_source": False} == Search().source(False).to_dict()
assert {
"_source": ["f1", "f2"]
} == Search().source(includes=["foo.bar.*"],
excludes=["foo.one"]).source(["f1", "f2"]).to_dict()
def test_response(self):
r = Response(
{
"took": 42,
"timed_out": False,
"_shards": {
"total": 10,
"successful": 10,
"skipped": 0,
"failed": 0
},
"hits": {
"total": {
"value": 34,
"relation": "eq"
},
"max_score":
0.0,
"hits": [
{
"_index": "my_index_01",
"_type": "_doc",
"_id": "1",
"_score": 1.0,
"_source": {
"field_23": 1
},
},
{
"_index": "my_index_01",
"_type": "_doc",
"_id": "2",
"_score": 0.2,
"_source": {
"field_23": 2
},
},
],
},
},
Search(),
)
self.assertEqual(r.took, 42)
self.assertEqual(r.timed_out, False)
self.assertEqual(r.success, True)
self.assertEqual(r._shards, {
"total": 10,
"successful": 10,
"skipped": 0,
"failed": 0
})
self.assertIsInstance(r.hits, Hits)
self.assertEqual(
r.__repr__(),
"<Response> took 42ms, success: True, total result 34, contains 2 hits",
)
def search(
self, nested_autocorrect: bool = True, repr_auto_execute: bool = True
) -> Search:
return Search(
using=self.client,
mappings=self.mappings,
index=self.name,
nested_autocorrect=nested_autocorrect,
repr_auto_execute=repr_auto_execute,
)
def test_expand__to_dot_is_respected(self):
s = Search().query("match", a__b=42, _expand__to_dot=False)
self.assertEqual({"query": {
"match": {
"a__b": {
"query": 42
}
}
}}, s.to_dict())
def test_search_query_combines_query():
s = Search()
s2 = s.query("match", f=42)
assert s2._query.to_dict() == Query(Match(f=42)).to_dict()
assert s._query.to_dict() is None
s3 = s2.query("match", f=43)
assert s2._query.to_dict() == Query(Match(f=42)).to_dict()
assert ordered(s3._query.to_dict()) == ordered(
Query(Bool(must=[Match(f=42), Match(f=43)])).to_dict())
def test_update_from_dict(self):
s = Search()
s.update_from_dict({"indices_boost": [{"important-documents": 2}]})
s.update_from_dict({"_source": ["id", "name"]})
assert {
"indices_boost": [{
"important-documents": 2
}],
"_source": ["id", "name"],
} == s.to_dict()
def test_repr_execution(client_search):
client_search.return_value = {
"took": 42,
"timed_out": False,
"_shards": {
"total": 10,
"successful": 10,
"skipped": 0,
"failed": 0
},
"hits": {
"total": {
"value": 34,
"relation": "eq"
},
"max_score":
0.0,
"hits": [
{
"_index": "my_index_01",
"_type": "_doc",
"_id": "1",
"_score": 1.0,
"_source": {
"field_23": 1
},
},
{
"_index": "my_index_01",
"_type": "_doc",
"_id": "2",
"_score": 0.2,
"_source": {
"field_23": 2
},
},
],
},
}
s = Search(using=Elasticsearch(hosts=["..."]),
index="yolo",
repr_auto_execute=True)
s.size(2).__repr__()
client_search.assert_called_once()
client_search.assert_any_call(body={"size": 2}, index=["yolo"])
client_search.reset_mock()
s.size(2)._repr_html_()
client_search.assert_called_once()
client_search.assert_any_call(body={"size": 2}, index=["yolo"])
def test_search_query_combines_query(self):
s = Search()
s2 = s.query("match", f=42)
self.assertEqual(s2._query.to_dict(), Query(Match(f=42)).to_dict())
self.assertEqual(s._query.to_dict(), None)
s3 = s2.query("match", f=43)
self.assertEqual(s2._query.to_dict(), Query(Match(f=42)).to_dict())
self.assertEqual(
ordered(s3._query.to_dict()),
ordered(Query(Bool(must=[Match(f=42), Match(f=43)])).to_dict()),
)
def test_search_column_selection():
mappings = Mappings(properties={
"col1": {
"type": "keyword"
},
"col2": {
"type": "integer"
}
})
assert Search(mappings=mappings)[["col1", "col2"]].to_dict() == {
"_source": {
"includes": ["col1", "col2"]
}
}
def test_suggest(self):
s = Search()
s = s.suggest("my_suggestion", "pyhton", term={"field": "title"})
assert {
"suggest": {
"my_suggestion": {
"term": {
"field": "title"
},
"text": "pyhton"
}
}
} == s.to_dict()
def test_grouping_agg(self):
my_agg = Aggs(sample.EXPECTED_AGG_QUERY, mappings=MAPPINGS)
agg_response = Aggregations(data=sample.ES_AGG_RESPONSE,
_search=Search().aggs(my_agg))
# none provided
self.assertIsNone(agg_response._grouping_agg()[0])
# fake provided
with self.assertRaises(KeyError):
agg_response._grouping_agg("yolo")
# not bucket provided
with self.assertRaises(ValueError):
agg_response._grouping_agg("avg_f1_micro")
# real provided
self.assertEqual(
agg_response._grouping_agg("global_metrics.field.name")[0],
"global_metrics.field.name",
)
def test_aggs_allow_two_metric():
s = Search()
s = s.aggs({"a": Max(field="a"), "b": Max(field="b")})
assert s.to_dict() == {
"aggs": {
"a": {
"max": {
"field": "a"
}
},
"b": {
"max": {
"field": "b"
}
}
}
}
def test_exclude():
s = Search()
s = s.exclude("match", title="python")
assert {
"query": {
"bool": {
"filter": [{
"bool": {
"must_not": [{
"match": {
"title": {
"query": "python"
}
}
}]
}
}]
}
}
} == s.to_dict()
def test_aggs_allow_two_metric(self):
s = Search()
s = s.aggs([Aggs("a", "max", field="a"), Aggs("b", "max", field="b")])
self.assertEqual(
s.to_dict(),
{
"aggs": {
"a": {
"max": {
"field": "a"
}
},
"b": {
"max": {
"field": "b"
}
}
}
},
)
def test_exclude(self):
s = Search()
s = s.exclude("match", title="python")
self.assertEqual(
{
"query": {
"bool": {
"filter": [{
"bool": {
"must_not": [{
"match": {
"title": {
"query": "python"
}
}
}]
}
}]
}
}
},
s.to_dict(),
)
def test_parse_as_dataframe(self):
my_agg = Aggs(sample.EXPECTED_AGG_QUERY, mappings=MAPPINGS)
df = Aggregations(data=sample.ES_AGG_RESPONSE,
_search=Search().aggs(my_agg)).to_dataframe(
grouped_by="global_metrics.field.name")
self.assertIsInstance(df, pd.DataFrame)
self.assertEqual(set(df.index.names),
{"classification_type", "global_metrics.field.name"})
self.assertEqual(set(df.columns),
{"avg_f1_micro", "avg_nb_classes", "doc_count"})
self.assertEqual(
df.to_dict(orient="index"),
{
("multiclass", "gpc"): {
"avg_f1_micro": 0.93,
"avg_nb_classes": 211.12,
"doc_count": 198,
},
("multiclass", "kind"): {
"avg_f1_micro": 0.89,
"avg_nb_classes": 206.5,
"doc_count": 370,
},
("multilabel", "ispracticecompatible"): {
"avg_f1_micro": 0.72,
"avg_nb_classes": 18.71,
"doc_count": 128,
},
("multilabel", "preservationmethods"): {
"avg_f1_micro": 0.8,
"avg_nb_classes": 9.97,
"doc_count": 76,
},
},
)
def search(self):
return Search(using=self.client,
mapping=self._mapping,
index=self.name)
def test_query_always_returns_search(self):
s = Search()
self.assertIsInstance(s.query("match", f=42), Search)
def test_source_on_clear(self):
assert ({} == Search().source(includes=["foo.bar.*"]).source(
includes=None, excludes=None).to_dict())
def test_complex_example(self):
s = (Search().query(
"match", title="python").must_not("match", title="ruby").should(
Query("term", category="meetup"),
Query("term", category="conference")).post_filter(
"terms", tags=["prague", "czech"]).script_fields(
more_attendees="doc['attendees'].value + 42").aggs(
"per_country", "terms", field="country").aggs(
"avg_attendees", "avg",
field="attendees").bool(
minimum_should_match=2).highlight_options(
order="score").highlight(
"title", "body", fragment_size=50))
self.assertEqual(
ordered({
"aggs": {
"per_country": {
"aggs": {
"avg_attendees": {
"avg": {
"field": "attendees"
}
}
},
"terms": {
"field": "country"
},
}
},
"highlight": {
"fields": {
"body": {
"fragment_size": 50
},
"title": {
"fragment_size": 50
},
},
"order": "score",
},
"post_filter": {
"terms": {
"tags": ["prague", "czech"]
}
},
"query": {
"bool": {
"minimum_should_match":
2,
"must": [{
"match": {
"title": {
"query": "python"
}
}
}],
"must_not": [{
"match": {
"title": {
"query": "ruby"
}
}
}],
"should": [
{
"term": {
"category": {
"value": "conference"
}
}
},
{
"term": {
"category": {
"value": "meetup"
}
}
},
],
}
},
"script_fields": {
"more_attendees": {
"script": "doc['attendees'].value + 42"
}
},
}),
ordered(s.to_dict()),
)
