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_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_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_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_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_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_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 search(self, nested_autocorrect=True):
return Search(
using=self.client,
mapping=self._mapping,
index=self.name,
nested_autocorrect=nested_autocorrect,
)
def test_suggest_accepts_global_text(self):
s = Search.from_dict({
"suggest": {
"text": "the amsterdma meetpu",
"my-suggest-1": {
"term": {
"field": "title"
}
},
"my-suggest-2": {
"text": "other",
"term": {
"field": "body"
}
},
}
})
assert {
"suggest": {
"my-suggest-1": {
"term": {
"field": "title"
},
"text": "the amsterdma meetpu",
},
"my-suggest-2": {
"term": {
"field": "body"
},
"text": "other"
},
}
} == s.to_dict()
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_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_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_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 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 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_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_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_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_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_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_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_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_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 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_from_dict_doesnt_need_query(self):
s = Search.from_dict({"size": 5})
assert {"size": 5} == s.to_dict()
def search(self):
return Search(using=self.client,
mapping=self._mapping,
index=self.name)
def test_reverse(self):
d = {
"aggs": {
"per_country": {
"aggs": {
"avg_attendees": {
"avg": {
"field": "attendees"
}
}
},
"terms": {
"field": "country"
},
}
},
"highlight": {
"fields": {
"title": {
"fragment_size": 50
}
},
"order": "score"
},
"post_filter": {
"bool": {
"must": [{
"terms": {
"tags": ["prague", "czech"]
}
}]
}
},
"query": {
"bool": {
"filter": [{
"bool": {
"should": [
{
"term": {
"category": {
"value": "meetup"
}
}
},
{
"term": {
"category": {
"value": "conference"
}
}
},
]
}
}],
"must": [{
"bool": {
"minimum_should_match": 2,
"must": [{
"match": {
"title": {
"query": "python"
}
}
}],
"must_not": [{
"match": {
"title": {
"query": "ruby"
}
}
}],
}
}],
}
},
"script_fields": {
"more_attendees": {
"script": "doc['attendees'].value + 42"
}
},
"size": 5,
"sort": ["title", {
"category": {
"order": "desc"
}
}, "_score"],
"suggest": {
"my-title-suggestions-1": {
"term": {
"field": "title",
"size": 3
},
"text": "devloping distibutd saerch "
"engies",
}
},
}
d2 = deepcopy(d)
s = Search.from_dict(d)
# make sure we haven't modified anything in place
self.assertEqual(d, d2)
self.assertSearchEqual(d, s.to_dict())