Best Python code snippet using SeleniumLibrary
test_tasks.py
Source:test_tasks.py
...25 reindex_project,26 unindex_project,27)28from ...common.db.packaging import FileFactory, ProjectFactory, ReleaseFactory29def test_project_docs(db_session):30 projects = [ProjectFactory.create() for _ in range(2)]31 releases = {32 p: sorted(33 [ReleaseFactory.create(project=p) for _ in range(3)],34 key=lambda r: packaging.version.parse(r.version),35 reverse=True,36 )37 for p in projects38 }39 for p in projects:40 for r in releases[p]:41 r.files = [42 FileFactory.create(43 release=r,44 filename="{}-{}.tar.gz".format(p.name, r.version),45 python_version="source",46 )47 ]48 assert list(_project_docs(db_session)) == [49 {50 "_id": p.normalized_name,51 "_source": {52 "created": p.created,53 "name": p.name,54 "normalized_name": p.normalized_name,55 "version": [r.version for r in prs],56 "latest_version": first(prs, key=lambda r: not r.is_prerelease).version,57 "description": first(58 prs, key=lambda r: not r.is_prerelease59 ).description.raw,60 },61 }62 for p, prs in sorted(releases.items(), key=lambda x: x[0].id)63 ]64def test_single_project_doc(db_session):65 projects = [ProjectFactory.create() for _ in range(2)]66 releases = {67 p: sorted(68 [ReleaseFactory.create(project=p) for _ in range(3)],69 key=lambda r: packaging.version.parse(r.version),70 reverse=True,71 )72 for p in projects73 }74 for p in projects:75 for r in releases[p]:76 r.files = [77 FileFactory.create(78 release=r,79 filename="{}-{}.tar.gz".format(p.name, r.version),80 python_version="source",81 )82 ]83 assert list(_project_docs(db_session, project_name=projects[1].name)) == [84 {85 "_id": p.normalized_name,86 "_source": {87 "created": p.created,88 "name": p.name,89 "normalized_name": p.normalized_name,90 "version": [r.version for r in prs],91 "latest_version": first(prs, key=lambda r: not r.is_prerelease).version,92 "description": first(93 prs, key=lambda r: not r.is_prerelease94 ).description.raw,95 },96 }97 for p, prs in sorted(releases.items(), key=lambda x: x[0].name.lower())98 if p.name == projects[1].name99 ]100def test_project_docs_empty(db_session):101 projects = [ProjectFactory.create() for _ in range(2)]102 releases = {103 p: sorted(104 [ReleaseFactory.create(project=p) for _ in range(3)],105 key=lambda r: packaging.version.parse(r.version),106 reverse=True,107 )108 for p in projects109 }110 project_with_files = projects[0]111 for r in releases[project_with_files]:112 r.files = [113 FileFactory.create(114 release=r,115 filename="{}-{}.tar.gz".format(project_with_files.name, r.version),116 python_version="source",117 )118 ]119 assert list(_project_docs(db_session)) == [120 {121 "_id": p.normalized_name,122 "_source": {123 "created": p.created,124 "name": p.name,125 "normalized_name": p.normalized_name,126 "version": [r.version for r in prs],127 "latest_version": first(prs, key=lambda r: not r.is_prerelease).version,128 "description": first(129 prs, key=lambda r: not r.is_prerelease130 ).description.raw,131 },132 }133 for p, prs in sorted(releases.items(), key=lambda x: x[0].id)134 if p == project_with_files135 ]136class FakeESIndices:137 def __init__(self):138 self.indices = {}139 self.aliases = {}140 self.put_settings = pretend.call_recorder(lambda *a, **kw: None)141 self.delete = pretend.call_recorder(lambda *a, **kw: None)142 self.create = pretend.call_recorder(lambda *a, **kw: None)143 def exists_alias(self, name):144 return name in self.aliases145 def get_alias(self, name):146 return self.aliases[name]147 def put_alias(self, name, index):148 self.aliases.setdefault(name, []).append(index)149 def remove_alias(self, name, alias):150 self.aliases[name] = [n for n in self.aliases[name] if n != alias]151 if not self.aliases[name]:152 del self.aliases[name]153 def update_aliases(self, body):154 for items in body["actions"]:155 for action, values in items.items():156 if action == "add":157 self.put_alias(values["alias"], values["index"])158 elif action == "remove":159 self.remove_alias(values["alias"], values["index"])160 else:161 raise ValueError("Unknown action: {!r}.".format(action))162class FakeESClient:163 def __init__(self):164 self.indices = FakeESIndices()165class NotLock:166 def __init__(*a, **kw):167 pass168 def acquire(self):169 return True170 def release(self):171 return True172class TestSearchLock:173 def test_success(self):174 lock_stub = pretend.stub(acquire=pretend.call_recorder(lambda: True))175 r = pretend.stub(lock=lambda *a, **kw: lock_stub)176 test_lock = SearchLock(r)177 test_lock.__enter__()178 assert lock_stub.acquire.calls == [pretend.call()]179 def test_failure(self):180 lock_stub = pretend.stub(acquire=pretend.call_recorder(lambda: False))181 r = pretend.stub(lock=lambda *a, **kw: lock_stub)182 test_lock = SearchLock(r)183 with pytest.raises(redis.exceptions.LockError):184 test_lock.__enter__()185 assert lock_stub.acquire.calls == [pretend.call()]186class TestReindex:187 def test_fails_when_raising(self, db_request, monkeypatch):188 docs = pretend.stub()189 def project_docs(db):190 return docs191 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)192 task = pretend.stub()193 es_client = FakeESClient()194 db_request.registry.update({"elasticsearch.index": "warehouse"})195 db_request.registry.settings = {196 "elasticsearch.url": "http://some.url",197 "celery.scheduler_url": "redis://redis:6379/0",198 }199 monkeypatch.setattr(200 warehouse.search.tasks.elasticsearch,201 "Elasticsearch",202 lambda *a, **kw: es_client,203 )204 class TestError(Exception):205 pass206 def parallel_bulk(client, iterable, index=None):207 assert client is es_client208 assert iterable is docs209 assert index == "warehouse-cbcbcbcbcb"210 raise TestError211 monkeypatch.setattr(212 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)213 )214 monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)215 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)216 with pytest.raises(TestError):217 reindex(task, db_request)218 assert es_client.indices.delete.calls == [219 pretend.call(index="warehouse-cbcbcbcbcb")220 ]221 assert es_client.indices.put_settings.calls == []222 def test_retry_on_lock(self, db_request, monkeypatch):223 task = pretend.stub(224 retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))225 )226 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}227 le = redis.exceptions.LockError()228 monkeypatch.setattr(229 redis.StrictRedis,230 "from_url",231 lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),232 )233 with pytest.raises(celery.exceptions.Retry):234 reindex(task, db_request)235 assert task.retry.calls == [pretend.call(countdown=60, exc=le)]236 def test_successfully_indexes_and_adds_new(self, db_request, monkeypatch):237 docs = pretend.stub()238 def project_docs(db):239 return docs240 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)241 task = pretend.stub()242 es_client = FakeESClient()243 db_request.registry.update(244 {"elasticsearch.index": "warehouse", "elasticsearch.shards": 42}245 )246 db_request.registry.settings = {247 "elasticsearch.url": "http://some.url",248 "celery.scheduler_url": "redis://redis:6379/0",249 }250 monkeypatch.setattr(251 warehouse.search.tasks.elasticsearch,252 "Elasticsearch",253 lambda *a, **kw: es_client,254 )255 monkeypatch.setattr(256 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)257 )258 parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])259 monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)260 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)261 reindex(task, db_request)262 assert parallel_bulk.calls == [263 pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")264 ]265 assert es_client.indices.create.calls == [266 pretend.call(267 body={268 "settings": {269 "number_of_shards": 42,270 "number_of_replicas": 0,271 "refresh_interval": "-1",272 }273 },274 wait_for_active_shards=42,275 index="warehouse-cbcbcbcbcb",276 )277 ]278 assert es_client.indices.delete.calls == []279 assert es_client.indices.aliases == {"warehouse": ["warehouse-cbcbcbcbcb"]}280 assert es_client.indices.put_settings.calls == [281 pretend.call(282 index="warehouse-cbcbcbcbcb",283 body={"index": {"number_of_replicas": 0, "refresh_interval": "1s"}},284 )285 ]286 def test_successfully_indexes_and_replaces(self, db_request, monkeypatch):287 docs = pretend.stub()288 task = pretend.stub()289 def project_docs(db):290 return docs291 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)292 es_client = FakeESClient()293 es_client.indices.indices["warehouse-aaaaaaaaaa"] = None294 es_client.indices.aliases["warehouse"] = ["warehouse-aaaaaaaaaa"]295 db_engine = pretend.stub()296 db_request.registry.update(297 {298 "elasticsearch.index": "warehouse",299 "elasticsearch.shards": 42,300 "sqlalchemy.engine": db_engine,301 }302 )303 db_request.registry.settings = {304 "elasticsearch.url": "http://some.url",305 "celery.scheduler_url": "redis://redis:6379/0",306 }307 monkeypatch.setattr(308 warehouse.search.tasks.elasticsearch,309 "Elasticsearch",310 lambda *a, **kw: es_client,311 )312 monkeypatch.setattr(313 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)314 )315 parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])316 monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)317 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)318 reindex(task, db_request)319 assert parallel_bulk.calls == [320 pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")321 ]322 assert es_client.indices.create.calls == [323 pretend.call(324 body={325 "settings": {326 "number_of_shards": 42,327 "number_of_replicas": 0,328 "refresh_interval": "-1",329 }330 },331 wait_for_active_shards=42,332 index="warehouse-cbcbcbcbcb",333 )334 ]335 assert es_client.indices.delete.calls == [pretend.call("warehouse-aaaaaaaaaa")]336 assert es_client.indices.aliases == {"warehouse": ["warehouse-cbcbcbcbcb"]}337 assert es_client.indices.put_settings.calls == [338 pretend.call(339 index="warehouse-cbcbcbcbcb",340 body={"index": {"number_of_replicas": 0, "refresh_interval": "1s"}},341 )342 ]343 def test_client_aws(self, db_request, monkeypatch):344 docs = pretend.stub()345 def project_docs(db):346 return docs347 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)348 task = pretend.stub()349 aws4auth_stub = pretend.stub()350 aws4auth = pretend.call_recorder(lambda *a, **kw: aws4auth_stub)351 es_client = FakeESClient()352 es_client_init = pretend.call_recorder(lambda *a, **kw: es_client)353 db_request.registry.update(354 {"elasticsearch.index": "warehouse", "elasticsearch.shards": 42}355 )356 db_request.registry.settings = {357 "aws.key_id": "AAAAAAAAAAAAAAAAAA",358 "aws.secret_key": "deadbeefdeadbeefdeadbeef",359 "elasticsearch.url": "https://some.url?aws_auth=1®ion=us-east-2",360 "celery.scheduler_url": "redis://redis:6379/0",361 }362 monkeypatch.setattr(363 warehouse.search.tasks.requests_aws4auth, "AWS4Auth", aws4auth364 )365 monkeypatch.setattr(366 warehouse.search.tasks.elasticsearch, "Elasticsearch", es_client_init367 )368 monkeypatch.setattr(369 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)370 )371 parallel_bulk = pretend.call_recorder(lambda client, iterable, index: [None])372 monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)373 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)374 reindex(task, db_request)375 assert len(es_client_init.calls) == 1376 assert es_client_init.calls[0].kwargs["hosts"] == ["https://some.url"]377 assert es_client_init.calls[0].kwargs["timeout"] == 30378 assert es_client_init.calls[0].kwargs["retry_on_timeout"] is True379 assert (380 es_client_init.calls[0].kwargs["connection_class"]381 == elasticsearch.connection.http_requests.RequestsHttpConnection382 )383 assert es_client_init.calls[0].kwargs["http_auth"] == aws4auth_stub384 assert aws4auth.calls == [385 pretend.call(386 "AAAAAAAAAAAAAAAAAA", "deadbeefdeadbeefdeadbeef", "us-east-2", "es"387 )388 ]389 assert parallel_bulk.calls == [390 pretend.call(es_client, docs, index="warehouse-cbcbcbcbcb")391 ]392 assert es_client.indices.create.calls == [393 pretend.call(394 body={395 "settings": {396 "number_of_shards": 42,397 "number_of_replicas": 0,398 "refresh_interval": "-1",399 }400 },401 wait_for_active_shards=42,402 index="warehouse-cbcbcbcbcb",403 )404 ]405 assert es_client.indices.delete.calls == []406 assert es_client.indices.aliases == {"warehouse": ["warehouse-cbcbcbcbcb"]}407 assert es_client.indices.put_settings.calls == [408 pretend.call(409 index="warehouse-cbcbcbcbcb",410 body={"index": {"number_of_replicas": 0, "refresh_interval": "1s"}},411 )412 ]413class TestPartialReindex:414 def test_reindex_fails_when_raising(self, db_request, monkeypatch):415 docs = pretend.stub()416 task = pretend.stub()417 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}418 def project_docs(db, project_name=None):419 return docs420 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)421 es_client = FakeESClient()422 db_request.registry.update(423 {"elasticsearch.client": es_client, "elasticsearch.index": "warehouse"}424 )425 class TestError(Exception):426 pass427 def parallel_bulk(client, iterable, index=None):428 assert client is es_client429 assert iterable is docs430 raise TestError431 monkeypatch.setattr(warehouse.search.tasks, "parallel_bulk", parallel_bulk)432 monkeypatch.setattr(433 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)434 )435 with pytest.raises(TestError):436 reindex_project(task, db_request, "foo")437 assert es_client.indices.put_settings.calls == []438 def test_unindex_fails_when_raising(self, db_request, monkeypatch):439 task = pretend.stub()440 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}441 class TestError(Exception):442 pass443 es_client = FakeESClient()444 es_client.delete = pretend.raiser(TestError)445 monkeypatch.setattr(446 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)447 )448 db_request.registry.update(449 {"elasticsearch.client": es_client, "elasticsearch.index": "warehouse"}450 )451 with pytest.raises(TestError):452 unindex_project(task, db_request, "foo")453 def test_unindex_accepts_defeat(self, db_request, monkeypatch):454 task = pretend.stub()455 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}456 es_client = FakeESClient()457 es_client.delete = pretend.call_recorder(458 pretend.raiser(elasticsearch.exceptions.NotFoundError)459 )460 monkeypatch.setattr(461 redis.StrictRedis, "from_url", lambda *a, **kw: pretend.stub(lock=NotLock)462 )463 db_request.registry.update(464 {"elasticsearch.client": es_client, "elasticsearch.index": "warehouse"}465 )466 unindex_project(task, db_request, "foo")467 assert es_client.delete.calls == [pretend.call(index="warehouse", id="foo")]468 def test_unindex_retry_on_lock(self, db_request, monkeypatch):469 task = pretend.stub(470 retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))471 )472 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}473 le = redis.exceptions.LockError()474 monkeypatch.setattr(475 redis.StrictRedis,476 "from_url",477 lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),478 )479 with pytest.raises(celery.exceptions.Retry):480 unindex_project(task, db_request, "foo")481 assert task.retry.calls == [pretend.call(countdown=60, exc=le)]482 def test_reindex_retry_on_lock(self, db_request, monkeypatch):483 task = pretend.stub(484 retry=pretend.call_recorder(pretend.raiser(celery.exceptions.Retry))485 )486 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}487 le = redis.exceptions.LockError()488 monkeypatch.setattr(489 redis.StrictRedis,490 "from_url",491 lambda *a, **kw: pretend.stub(lock=pretend.raiser(le)),492 )493 with pytest.raises(celery.exceptions.Retry):494 reindex_project(task, db_request, "foo")495 assert task.retry.calls == [pretend.call(countdown=60, exc=le)]496 def test_successfully_indexes(self, db_request, monkeypatch):497 docs = pretend.stub()498 task = pretend.stub()499 db_request.registry.settings = {"celery.scheduler_url": "redis://redis:6379/0"}500 def project_docs(db, project_name=None):501 return docs502 monkeypatch.setattr(warehouse.search.tasks, "_project_docs", project_docs)503 es_client = FakeESClient()504 es_client.indices.indices["warehouse-aaaaaaaaaa"] = None505 es_client.indices.aliases["warehouse"] = ["warehouse-aaaaaaaaaa"]506 db_engine = pretend.stub()507 db_request.registry.update(508 {509 "elasticsearch.client": es_client,510 "elasticsearch.index": "warehouse",511 "elasticsearch.shards": 42,512 "sqlalchemy.engine": db_engine,513 }514 )...
test_reindex.py
Source:test_reindex.py
...15from first import first16import warehouse.cli.search.reindex17from warehouse.cli.search.reindex import reindex, _project_docs18from ....common.db.packaging import ProjectFactory, ReleaseFactory19def test_project_docs(db_session):20 projects = [ProjectFactory.create() for _ in range(2)]21 releases = {22 p: sorted(23 [ReleaseFactory.create(project=p) for _ in range(3)],24 key=lambda r: packaging.version.parse(r.version),25 reverse=True,26 )27 for p in projects28 }29 assert list(_project_docs(db_session)) == [30 {31 "_id": p.normalized_name,32 "_type": "project",33 "_source": {34 "created": p.created,35 "name": p.name,36 "normalized_name": p.normalized_name,37 "version": [r.version for r in prs],38 "latest_version": first(39 prs,40 key=lambda r: not r.is_prerelease,41 ).version,42 },43 }44 for p, prs in sorted(releases.items(), key=lambda x: x[0].name.lower())45 ]46class FakeESIndices:47 def __init__(self):48 self.indices = {}49 self.aliases = {}50 self.put_settings = pretend.call_recorder(lambda *a, **kw: None)51 self.forcemerge = pretend.call_recorder(lambda *a, **kw: None)52 self.delete = pretend.call_recorder(lambda *a, **kw: None)53 self.create = pretend.call_recorder(lambda *a, **kw: None)54 def exists_alias(self, name):55 return name in self.aliases56 def get_alias(self, name):57 return self.aliases[name]58 def put_alias(self, name, index):59 self.aliases.setdefault(name, []).append(index)60 def remove_alias(self, name, alias):61 self.aliases[name] = [n for n in self.aliases[name] if n != alias]62 if not self.aliases[name]:63 del self.aliases[name]64 def update_aliases(self, body):65 for items in body["actions"]:66 for action, values in items.items():67 if action == "add":68 self.put_alias(values["alias"], values["index"])69 elif action == "remove":70 self.remove_alias(values["alias"], values["index"])71 else:72 raise ValueError("Unknown action: {!r}.".format(action))73class FakeESClient:74 def __init__(self):75 self.indices = FakeESIndices()76class TestReindex:77 def test_fails_when_raising(self, monkeypatch, cli):78 sess_obj = pretend.stub(79 execute=pretend.call_recorder(lambda q: None),80 rollback=pretend.call_recorder(lambda: None),81 close=pretend.call_recorder(lambda: None),82 )83 sess_cls = pretend.call_recorder(lambda bind: sess_obj)84 monkeypatch.setattr(warehouse.cli.search.reindex, "Session", sess_cls)85 docs = pretend.stub()86 def project_docs(db):87 return docs88 monkeypatch.setattr(89 warehouse.cli.search.reindex,90 "_project_docs",91 project_docs,92 )93 es_client = FakeESClient()94 db_engine = pretend.stub()95 config = pretend.stub(96 registry={97 "elasticsearch.client": es_client,98 "elasticsearch.index": "warehouse",99 "sqlalchemy.engine": db_engine,100 },101 )102 class TestException(Exception):103 pass104 def parallel_bulk(client, iterable):105 assert client is es_client106 assert iterable is docs107 raise TestException108 monkeypatch.setattr(109 warehouse.cli.search.reindex, "parallel_bulk", parallel_bulk)110 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)111 result = cli.invoke(reindex, obj=config)112 assert result.exit_code == -1113 assert isinstance(result.exception, TestException)114 assert sess_cls.calls == [pretend.call(bind=db_engine)]115 assert sess_obj.execute.calls == [116 pretend.call("SET statement_timeout = '600s'"),117 ]118 assert sess_obj.rollback.calls == [pretend.call()]119 assert sess_obj.close.calls == [pretend.call()]120 assert es_client.indices.delete.calls == [121 pretend.call(index='warehouse-cbcbcbcbcb'),122 ]123 assert es_client.indices.put_settings.calls == []124 assert es_client.indices.forcemerge.calls == []125 def test_successfully_indexes_and_adds_new(self, monkeypatch, cli):126 sess_obj = pretend.stub(127 execute=pretend.call_recorder(lambda q: None),128 rollback=pretend.call_recorder(lambda: None),129 close=pretend.call_recorder(lambda: None),130 )131 sess_cls = pretend.call_recorder(lambda bind: sess_obj)132 monkeypatch.setattr(warehouse.cli.search.reindex, "Session", sess_cls)133 docs = pretend.stub()134 def project_docs(db):135 return docs136 monkeypatch.setattr(137 warehouse.cli.search.reindex,138 "_project_docs",139 project_docs,140 )141 es_client = FakeESClient()142 db_engine = pretend.stub()143 config = pretend.stub(144 registry={145 "elasticsearch.client": es_client,146 "elasticsearch.index": "warehouse",147 "elasticsearch.shards": 42,148 "sqlalchemy.engine": db_engine,149 },150 )151 parallel_bulk = pretend.call_recorder(lambda client, iterable: [None])152 monkeypatch.setattr(153 warehouse.cli.search.reindex, "parallel_bulk", parallel_bulk)154 monkeypatch.setattr(os, "urandom", lambda n: b"\xcb" * n)155 result = cli.invoke(reindex, obj=config)156 assert result.exit_code == 0157 assert sess_cls.calls == [pretend.call(bind=db_engine)]158 assert sess_obj.execute.calls == [159 pretend.call("SET statement_timeout = '600s'"),160 ]161 assert parallel_bulk.calls == [pretend.call(es_client, docs)]162 assert sess_obj.rollback.calls == [pretend.call()]163 assert sess_obj.close.calls == [pretend.call()]164 assert es_client.indices.create.calls == [165 pretend.call(166 body={167 'settings': {168 'number_of_shards': 42,169 'number_of_replicas': 0,170 'refresh_interval': '-1',171 }172 },173 wait_for_active_shards=42,174 index='warehouse-cbcbcbcbcb',175 )176 ]177 assert es_client.indices.delete.calls == []178 assert es_client.indices.aliases == {179 "warehouse": ["warehouse-cbcbcbcbcb"],180 }181 assert es_client.indices.put_settings.calls == [182 pretend.call(183 index='warehouse-cbcbcbcbcb',184 body={185 'index': {186 'number_of_replicas': 0,187 'refresh_interval': '1s',188 },189 },190 )191 ]192 assert es_client.indices.forcemerge.calls == [193 pretend.call(index='warehouse-cbcbcbcbcb')194 ]195 def test_successfully_indexes_and_replaces(self, monkeypatch, cli):196 sess_obj = pretend.stub(197 execute=pretend.call_recorder(lambda q: None),198 rollback=pretend.call_recorder(lambda: None),199 close=pretend.call_recorder(lambda: None),200 )201 sess_cls = pretend.call_recorder(lambda bind: sess_obj)202 monkeypatch.setattr(warehouse.cli.search.reindex, "Session", sess_cls)203 docs = pretend.stub()204 def project_docs(db):205 return docs206 monkeypatch.setattr(207 warehouse.cli.search.reindex,208 "_project_docs",209 project_docs,210 )211 es_client = FakeESClient()212 es_client.indices.indices["warehouse-aaaaaaaaaa"] = None213 es_client.indices.aliases["warehouse"] = ["warehouse-aaaaaaaaaa"]214 db_engine = pretend.stub()215 config = pretend.stub(216 registry={217 "elasticsearch.client": es_client,218 "elasticsearch.index": "warehouse",...
managers.py
Source:managers.py
1import json2from django.db import models3from django.db.models import Prefetch4from .constants import NAMED_ENTITY_RECOGNITION_VALUE, RELATION_EXTRACTION_VALUE, SENTIMENT_ANALYSIS_VALUE5from .utils import convert_task_name_to_annotation_type, convert_explanation_int_to_explanation_type, SPACY_WRAPPER6class AnnotationManager(models.Manager):7 def get_annotations_for_export(self, user_id, task_name, explanation_int):8 queryset = self.get_queryset()9 annotation_type = convert_task_name_to_annotation_type(task_name)10 if explanation_int > 1:11 explanation_type = convert_explanation_int_to_explanation_type(explanation_int)12 annotation_queryset = queryset.filter(user_id=user_id) \13 .select_related(annotation_type, "label") \14 .prefetch_related(explanation_type)15 else:16 annotation_queryset = queryset.filter(user_id=user_id) \17 .select_related(annotation_type, "label")18 19 return annotation_queryset20class DocumentManager(models.Manager):21 def _format_json(self, project_docs, user_id, explanation, task_name):22 dataset = {"data" : []}23 for doc in project_docs:24 annotations = []25 for a in doc.user_annotations:26 if not a.user_provided:27 output = {28 "annotation_id" : a.id,29 "label" : a.label.text,30 }31 extended_output = a.format_json_outputs(task_name)32 for key in extended_output:33 output[key] = extended_output[key]34 if explanation:35 output["explanations"] = a.format_explanations("json", doc.text)36 annotations.append(output)37 dataset["data"].append({38 "doc_id" : doc.id,39 "text" : doc.text,40 "annotations" : annotations,41 "user" : user_id,42 "metadata" : json.loads(doc.metadata)43 })44 return dataset45 46 def _format_csv_ner(self, project_docs, explanation):47 dataset = []48 header = ["document_id", "word", "label", "metadata"]49 if explanation:50 header.append("explanation")51 dataset.append(header)52 for i, doc in enumerate(project_docs):53 text = doc.text54 words = SPACY_WRAPPER.tokenize(text)55 if explanation:56 doc_rep = [[i+1, word.encode('utf-8'), 'O', "", ""] for word in words]57 else:58 doc_rep = [[i+1, word.encode('utf-8'), 'O', ""] for word in words]59 doc_rep[0][3] = doc.metadata60 startoff_map = {}61 start_off = 062 for word_index, tup in enumerate(doc_rep):63 startoff_map[start_off] = word_index64 start_off = start_off + len(tup[1]) + 165 for a in doc.user_annotations:66 start_offset = a.named_entity_annotation.start_offset67 end_offset = a.named_entity_annotation.end_offset68 if start_offset in startoff_map:69 doc_rep[startoff_map[start_offset]][2] = 'B-{}'.format(a.label.text)70 if explanation:71 explanations = a.format_explanations("csv", doc.text)72 doc_rep[startoff_map[start_offset]][4] = explanations73 for i in range(start_offset+1, end_offset):74 if i in startoff_map:75 doc_rep[startoff_map[i]][2] = 'I-{}'.format(a.label.text)76 77 for row in doc_rep:78 dataset.append(row)79 dataset.append("")80 81 dataset.pop()82 return dataset83 def _format_csv_re(self, project_docs, explanation):84 dataset = []85 header = ["document_id", "entity_1", "entity_2", "label", "text", "metadata"]86 if explanation:87 header.append("explanation")88 dataset.append(header)89 for i, doc in enumerate(project_docs):90 text = doc.text91 for j, a in enumerate(doc.user_annotations):92 if not a.user_provided:93 tmp_start_offset = a.relation_extraction_annotation.sbj_start_offset94 tmp_end_offset = a.relation_extraction_annotation.sbj_end_offset95 sbj_entity = text[tmp_start_offset:tmp_end_offset]96 97 tmp_start_offset = a.relation_extraction_annotation.obj_start_offset98 tmp_end_offset = a.relation_extraction_annotation.obj_end_offset99 obj_entity = text[tmp_start_offset:tmp_end_offset]100 metadata = doc.metadata if j == 0 else ""101 102 if explanation:103 explanations = a.format_explanations("csv", doc.text)104 dataset.append([i+1, sbj_entity.encode('utf-8'), obj_entity.encode('utf-8'), a.label.text, text, metadata, explanations])105 else:106 dataset.append([i+1, sbj_entity.encode('utf-8'), obj_entity.encode('utf-8'), a.label.text, text, metadata])107 108 dataset.append("")109 110 dataset.pop()111 return dataset112 def _format_csv_sa(self, project_docs, explanation):113 dataset = []114 header = ["document_id", "text", "label", "metadata"]115 if explanation:116 header.append("explanation")117 dataset.append(header)118 for i, doc in enumerate(project_docs):119 for j, a in enumerate(doc.user_annotations):120 metadata = doc.metadata if j ==0 else ""121 if explanation:122 explanations = a.format_explanations("csv", doc.text)123 dataset.append([i+1, doc.text.encode('utf-8'), a.label.text, metadata, explanations])124 else:125 dataset.append([i+1, doc.text.encode('utf-8'), a.label.text, metadata])126 127 return dataset128 def export_ner_project_user_documents(self, project_id, user_id, export_format, annotation_queryset, explanation=False):129 queryset = self.get_queryset()130 # https://docs.djangoproject.com/en/3.1/ref/models/querysets/#django.db.models.Prefetch131 project_docs = queryset.filter(project_id=project_id).prefetch_related(Prefetch(132 lookup="annotations",133 queryset=annotation_queryset,134 to_attr="user_annotations"135 ))136 if export_format == "csv":137 return self._format_csv_ner(project_docs, explanation)138 139 if export_format == "json":140 return self._format_json(project_docs, user_id, explanation, NAMED_ENTITY_RECOGNITION_VALUE)141 142 143 def export_re_project_user_documents(self, project_id, user_id, export_format, annotation_queryset, explanation=False):144 queryset = self.get_queryset()145 # https://docs.djangoproject.com/en/3.1/ref/models/querysets/#django.db.models.Prefetch146 project_docs = queryset.filter(project_id=project_id).prefetch_related(Prefetch(147 lookup="annotations",148 queryset=annotation_queryset,149 to_attr="user_annotations"150 ))151 if export_format == "csv":152 return self._format_csv_re(project_docs, explanation)153 154 if export_format == "json":155 return self._format_json(project_docs, user_id, explanation, RELATION_EXTRACTION_VALUE)156 157 def export_sa_project_user_documents(self, project_id, user_id, export_format, annotation_queryset, explanation=False):158 queryset = self.get_queryset()159 project_docs = queryset.filter(project_id=project_id).prefetch_related(Prefetch(160 lookup="annotations",161 queryset=annotation_queryset,162 to_attr="user_annotations"163 ))164 if export_format == "csv":165 return self._format_csv_sa(project_docs, explanation)166 167 if export_format == "json":168 return self._format_json(project_docs, user_id, explanation, SENTIMENT_ANALYSIS_VALUE)169 170 def export_project_user_documents(self, task_name, project_id, user_id, export_format, annotation_queryset, explanation_int):171 explanation_bool = explanation_int > 1172 if task_name == NAMED_ENTITY_RECOGNITION_VALUE:173 return self.export_ner_project_user_documents(project_id, user_id, export_format, annotation_queryset, explanation_bool)174 elif task_name == RELATION_EXTRACTION_VALUE:175 return self.export_re_project_user_documents(project_id, user_id, export_format, annotation_queryset, explanation_bool)176 elif task_name == SENTIMENT_ANALYSIS_VALUE:177 return self.export_sa_project_user_documents(project_id, user_id, export_format, annotation_queryset, explanation_bool)178 179 def get_project_docs_with_labels(self, project_id):180 queryset = self.get_queryset()181 project_docs = queryset.filter(project_id=project_id) \182 .prefetch_related("annotations__label", 183 "annotations__user")...
test_redactor.py
Source:test_redactor.py
1import argparse2import os3from project1 import redactor4# Testing creation of status file5def testStatusFile():6 # Calling the method to create testlog file7 redactor.statsFile('../test_project/testLog')8 # Verifying if the testlog file is created or not9 assert os.path.isfile('project_docs/test_project/testlog')10# Testing the method used to updated status log11def testUpdateStatusLog():12 # Message to be writted to status log13 message = 'Testing status updating method'14 # Updating the status log15 redactor.updateStatusLog(message, '../test_project/testLog')16 assert open(17 'project_docs/test_project/testlog').read().splitlines()[-1] == message # Verifying the updated message18# Testing the method used to write data to redacted files19def testRedactedDoc():20 # Data to writted in redacted file and it's name21 message = ('\u2588 Testing redactedDoc method \u2588', 'test')22 redactor.redactedDoc(message, 'test_project') # Calling method redactedDoc23 # Verifying if the file with .redacted extension id created or not24 assert os.path.isfile('project_docs/test_project/test.redacted')25 assert open(26 'project_docs/test_project/test.redacted').read().splitlines()[-1] == message[0] # Verifying the contents of the redacted file27# Testing the method to fetch or read data from the input files.28def testFetchDocs():29 parser = argparse.ArgumentParser() # Creating an argument parser object30 # Adding optinal argument --inputs31 parser.add_argument("--input", type=str, required=True,32 nargs="+", action="append")33 parser.add_argument("--stats", type=str, required=False,34 default="stdout") # Adding optinal argument --stats35 args = parser.parse_args(36 "--input ../tests/test.txt --stats ../test_project/testLog".split())37 # Calling the method to fetch data from documents38 global data39 data = redactor.fetchDocs(args)40 # Verifying if the the return type of method is list or not41 assert type(data) == list42 assert len(data[0][0]) > 1 # Verifying if the read file has data43 # Verifying the name of the file from which data is read44 assert data[0][1] == 'test'45# Testing the method used to redact Names46def testRedactNames():47 parser = argparse.ArgumentParser() # Creating an argument parser object48 parser.add_argument("--names", required=False, action="store_true")49 parser.add_argument("--stats", type=str, required=False,50 default="stdout") # Adding optinal argument --stats51 args = parser.parse_args(52 "--names --stats ../test_project/testLog".split())53 redactor.redactNames(data[0][0], args)54 # log file must be updated with this message as there are 9 names to be redacted in test document55 expected = 'Sucessfully redacted 9 Names'56 assert open(57 'project_docs/test_project/testLog').read().splitlines()[-1] == expected # Verifying the contents of the log file w.r.t expected58# Testing the method used to redact Phone Numbers59def testRedactPhones():60 parser = argparse.ArgumentParser() # Creating an argument parser object61 parser.add_argument("--phones", required=False, action="store_true")62 parser.add_argument("--stats", type=str, required=False,63 default="stdout") # Adding optinal argument --stats64 args = parser.parse_args(65 "--phones --stats ../test_project/testLog".split())66 redactor.redactPhone(data[0][0], args)67 # log file must be updated with this message as there are 8 numbers to be redacted in test document68 expected = 'Sucessfully redacted 8 phone numbers'69 assert open(70 'project_docs/test_project/testLog').read().splitlines()[-1] == expected # Verifying the contents of the log file w.r.t expected71# Testing the method used to redact gender related words72def testRedactGenders():73 parser = argparse.ArgumentParser() # Creating an argument parser object74 parser.add_argument("--genders", required=False, action="store_true")75 parser.add_argument("--stats", type=str, required=False,76 default="stdout") # Adding optinal argument --stats77 args = parser.parse_args(78 "--genders --stats ../test_project/testLog".split())79 redactor.redactGender(data[0][0], args)80 # log file must be updated with this message as there are 17 words related to gender in test document81 expected = 'Sucessfully redacted 17 words related to gender identity'82 assert open(83 'project_docs/test_project/testLog').read().splitlines()[-1] == expected # Verifying the contents of the log file w.r.t expected84# Testing the method used to redact dates85def testRedactDates():86 parser = argparse.ArgumentParser() # Creating an argument parser object87 parser.add_argument("--dates", required=False, action="store_true")88 parser.add_argument("--stats", type=str, required=False,89 default="stdout") # Adding optinal argument --stats90 args = parser.parse_args(91 "--dates --stats ../test_project/testLog".split())92 redactor.redactDates(data[0][0], args)93 # log file must be updated with this message as there are 8 dates in test document94 expected = 'Sucessfully redacted 8 Dates'95 assert open(96 'project_docs/test_project/testLog').read().splitlines()[-1] == expected # Verifying the contents of the log file w.r.t expected97# Testing the method used to redact sentences containing the words related to concept98def testRedactConcept():99 parser = argparse.ArgumentParser() # Creating an argument parser object100 # Adding optinal argument --concept101 parser.add_argument("--concept", type=str,102 required=False, nargs="+", action="append")103 parser.add_argument("--stats", type=str, required=False,104 default="stdout") # Adding optinal argument --stats105 args = parser.parse_args(106 "--concept shot --stats ../test_project/testLog".split())107 redactor.redactConcept(data[0][0], args)108 # log file must be updated with this message as there are 4 sentences in test document containing the words related to concept 'shot'109 expected = 'Sucessfully redacted 4 sentences where words related to concept are present'110 assert open(...
Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!