from __future__ import print_function from future import standard_library standard_library.install_aliases() from builtins import range from builtins import str from builtins import object from past.builtins import basestring import pytest import os try: import subprocess32 as subprocess except: import subprocess import sys import psutil import time import re import json import urllib.request, urllib.error, urllib.parse sys.path.append(os.getcwd()) from mrq.job import Job, queue_raw_jobs, queue_jobs from mrq.queue import Queue from mrq.config import get_config from mrq.utils import wait_for_net_service from mrq.context import connections, set_current_config set_current_config(get_config(sources=("env"))) os.system("rm -rf dump.rdb") PYTHON_BIN = "python" if os.environ.get("PYTHON_BIN"): PYTHON_BIN = os.environ["PYTHON_BIN"] class ProcessFixture(object): def __init__(self, request, cmdline=None, wait_port=None, quiet=False): self.request = request self.cmdline = cmdline self.process = None self.wait_port = wait_port self.quiet = quiet self.stopped = False self.started = False self.request.addfinalizer(self.stop) def start(self, cmdline=None, env=None, expected_children=0): self.stopped = False self.process_children = [] if not cmdline: cmdline = self.cmdline if env is None: env = {} # Kept from parent env for env_key in ["PATH", "GEVENT_LOOP", "VIRTUAL_ENV"]: if os.environ.get(env_key) and not env.get(env_key): env[env_key] = os.environ.get(env_key) if self.quiet: stdout = open(os.devnull, 'w') else: stdout = None self.cmdline = cmdline # print cmdline self.process = subprocess.Popen(re.split(r"\s+", cmdline) if isinstance(cmdline, basestring) else cmdline, shell=False, close_fds=True, env=env, cwd=os.getcwd(), stdout=stdout) if self.quiet: stdout.close() # Wait for children to start if expected_children > 0: psutil_process = psutil.Process(self.process.pid) while True: self.process_children = psutil_process.children(recursive=True) # print("Expecting %s children of pid %s, got %s" % ( # expected_children, self.process.pid, len(self.process_children)) # ) if len(self.process_children) >= expected_children: break time.sleep(0.1) if self.wait_port: wait_for_net_service("127.0.0.1", int(self.wait_port), poll_interval=0.01) self.started = True def stop(self, force=False, timeout=None, block=True, sig=15): # Call this only one time. if self.stopped and not force: return self.stopped = True self.started = False if self.process is not None and self.process.returncode is None: print("Sending signal %s to pid %s" % (sig, self.process.pid)) os.kill(self.process.pid, sig) # When sending a sigkill to the process, we also want to kill the # children in case of supervisord usage if sig == 9 and len(self.process_children) > 0: for c in self.process_children: c.send_signal(sig) if not block: return # Will raise TimeoutExpired self.process.wait(timeout=20) class WorkerFixture(ProcessFixture): def __init__(self, request, **kwargs): ProcessFixture.__init__(self, request, cmdline=kwargs.get("cmdline")) self.fixture_mongodb = kwargs["mongodb"] self.fixture_redis = kwargs["redis"] self.admin_port = kwargs.get("admin_port") self.started = False self.agent = False def start(self, flush=True, deps=True, trace=True, agent=False, **kwargs): self.started = True self.agent = agent if deps: self.start_deps(flush=flush) processes = 0 if self.agent: cmdline = "%s mrq/bin/mrq_agent.py %s" % (PYTHON_BIN, kwargs.get("flags", "")) else: m = re.search(r"--processes (\d+)", kwargs.get("flags", "")) if m: processes = int(m.group(1)) cmdline = "%s mrq/bin/mrq_worker.py --mongodb_logs_size 0 %s %s %s %s" % ( PYTHON_BIN, "--admin_port %s" % self.admin_port if (processes <= 1) else "", "--trace_io --trace_greenlets" if trace else "", kwargs.get("flags", ""), kwargs.get("queues", "high default low") ) env = kwargs.get("env") or {} env.setdefault("MRQ_MAX_LATENCY", "0.1") # Pass coverage-related environment variables for k, v in os.environ.items(): if k.startswith("COV_"): env[k] = v print(cmdline) ProcessFixture.start(self, cmdline=cmdline, env=env, expected_children=processes) if kwargs.get("block") is False: self.wait_for_idle() def start_deps(self, flush=True): self.fixture_mongodb.start() self.fixture_redis.start() # Will auto-connect connections.reset() self.mongodb_jobs = connections.mongodb_jobs self.mongodb_logs = connections.mongodb_logs self.redis = connections.redis if flush: self.fixture_mongodb.flush() self.fixture_redis.flush() def stop(self, deps=True, sig=2, **kwargs): if self.started: ProcessFixture.stop(self, sig=sig, **kwargs) if deps: self.stop_deps(**kwargs) def stop_deps(self, **kwargs): self.fixture_mongodb.stop(sig=2, **kwargs) self.fixture_redis.stop(sig=2, **kwargs) def wait_for_tasks_results(self, job_ids, block=True, accept_statuses=["success"]): if not block: return job_ids self.wait_for_idle() results = [] for job_id in job_ids: job = Job(job_id).wait(poll_interval=0.01) assert job.get("status") in accept_statuses, "Job had status %s, not in %s. Dump: %s" % ( job.get("status"), accept_statuses, job) results.append(job.get("result")) return results def send_raw_tasks(self, queue, params_list, start=True, block=True): if not self.started and start: self.start() queue_raw_jobs(queue, params_list) if block: self.wait_for_idle() def send_tasks(self, path, params_list, block=True, queue=None, accept_statuses=["success"], start=True): if not self.started and start: self.start() job_ids = queue_jobs(path, params_list, queue=queue) return self.wait_for_tasks_results(job_ids, block=block, accept_statuses=accept_statuses) def send_task(self, path, params, **kwargs): return self.send_tasks(path, [params], **kwargs)[0] def send_task_cli(self, path, params, queue=None, **kwargs): cli = [PYTHON_BIN, "mrq/bin/mrq_run.py", "--quiet"] if queue: cli += ["--queue", queue] cli += [path, json.dumps(params)] out = subprocess.check_output(cli).strip() if not queue: return json.loads(out.decode('utf-8')) return out def get_report(self, with_memory=False): wait_for_net_service("127.0.0.1", self.admin_port, poll_interval=0.01) f = urllib.request.urlopen("http://127.0.0.1:%s/report%s" % (self.admin_port, "_mem" if with_memory else "")) data = json.loads(f.read().decode('utf-8')) f.close() return data def wait_for_idle(self): if "--processes" in self.cmdline: print("Warning: wait_for_idle() doesn't support multiprocess workers yet") return False if self.agent: print("Warning: wait_for_idle() doesn't support agent yet") return False if self.process.returncode is not None: return True try: wait_for_net_service("127.0.0.1", self.admin_port, poll_interval=0.01) f = urllib.request.urlopen("http://127.0.0.1:%s/wait_for_idle" % self.admin_port) data = f.read().decode('utf-8') assert data == "idle" return True except Exception as e: print("Couldn't wait_for_idle: %s" % e) return False class RedisFixture(ProcessFixture): def flush(self): connections.redis.flushall() # Empty local known_queues cache too Queue.known_queues = {} class MongoFixture(ProcessFixture): def flush(self): for mongodb in (connections.mongodb_jobs, connections.mongodb_logs): if mongodb: for c in mongodb.collection_names(): if not c.startswith("system."): mongodb.drop_collection(c) class ApiFixture(ProcessFixture): def __init__(self, *args, **kwargs): ProcessFixture.__init__(self, *args, **kwargs) self.wait_port = 5555 self.host = "127.0.0.1" self.url = "http://%s:%s" % (self.host, self.wait_port) self.cmdline = "%s mrq/dashboard/app.py" % PYTHON_BIN def start(self, *args, **kwargs): kwargs["env"] = os.environ ProcessFixture.start(self, *args, **kwargs) def _login(self, user): import requests res = requests.post(self.url + "/api/auth/login", data={ "userEmail": user["email"], "userPassword": user["password"] }) assert res.status_code == 200 return { "cookies": res.cookies } def _request(self, method, path, key="local", user=None, assert_200=True, **kwargs): if not self.started: self.start() import requests if user: auth = self._login(user) kwargs["cookies"] = auth["cookies"] res = getattr(requests, method)(self.url + path, **kwargs) if assert_200 and res.status_code != 200: raise Exception(res.status_code) try: js = json.loads(res.text) except: print("Couldn't json parse", repr(res.text)[0:100]) js = None return js, res def GET(self, *args, **kwargs): return self._request("get", *args, **kwargs) def POST(self, *args, **kwargs): return self._request("post", *args, **kwargs) def DELETE(self, *args, **kwargs): return self._request("delete", *args, **kwargs) @pytest.fixture(scope="function") def httpstatic(request): return ProcessFixture(request, "/usr/sbin/nginx -c /app/tests/fixtures/httpstatic/nginx.conf", wait_port=8081) @pytest.fixture(scope="function") def mongodb(request): cmd = "mongod --smallfiles --noprealloc --nojournal" if os.environ.get("STACK_STARTED"): cmd = "sleep 1h" return MongoFixture(request, cmd, wait_port=27017, quiet=True) @pytest.fixture(scope="function") def mongodb_with_journal(request): cmd = "mongod --smallfiles --noprealloc" if os.environ.get("STACK_STARTED"): cmd = "sleep 1h" return MongoFixture(request, cmd, wait_port=27017, quiet=True) @pytest.fixture(scope="function") def redis(request): cmd = "redis-server" if os.environ.get("STACK_STARTED"): cmd = "sleep 1h" return RedisFixture(request, cmd, wait_port=6379, quiet=True) @pytest.fixture(scope="function") def worker(request, mongodb, redis): return WorkerFixture(request, mongodb=mongodb, redis=redis, admin_port=20020) @pytest.fixture(scope="function") def worker2(request, mongodb, redis): return WorkerFixture(request, mongodb=mongodb, redis=redis, admin_port=20021) @pytest.fixture(scope="function") def worker_mongodb_with_journal(request, mongodb_with_journal, redis): return WorkerFixture(request, mongodb=mongodb_with_journal, redis=redis, admin_port=20022) @pytest.fixture(scope="function") def api(request, mongodb, redis): return ApiFixture(request)