def test_env_priority(mocked_connections, monkeypatch):
"""Env vars should be used over defaults in connect"""
# Use monkeypatch to set environment variable only for this test
monkeypatch.setenv("LOG_LEVEL", "DUMMY_ENV_LOG_LEVEL")
monkeypatch.setenv("FLOWDB_PORT", "6969")
monkeypatch.setenv("FLOWDB_USER", "DUMMY_ENV_FLOWDB_USER")
monkeypatch.setenv("FLOWDB_PASS", "DUMMY_ENV_FLOWDB_PASS")
monkeypatch.setenv("FLOWDB_HOST", "DUMMY_ENV_FLOWDB_HOST")
monkeypatch.setenv("DB_CONNECTION_POOL_SIZE", "7777")
monkeypatch.setenv("DB_CONNECTION_POOL_OVERFLOW", "2020")
monkeypatch.setenv("REDIS_HOST", "DUMMY_ENV_REDIS_HOST")
monkeypatch.setenv("REDIS_PORT", "5050")
monkeypatch.setenv("REDIS_PASSWORD", "DUMMY_ENV_REDIS_PASSWORD")
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect()
core_init_logging_mock.assert_called_with("DUMMY_ENV_LOG_LEVEL")
core_init_Connection_mock.assert_called_with(
port=6969,
user="******",
password="******",
host="DUMMY_ENV_FLOWDB_HOST",
database="flowdb",
pool_size=7777,
overflow=2020,
)
core_init_StrictRedis_mock.assert_called_with(
host="DUMMY_ENV_REDIS_HOST",
port=5050,
password="******")
core_init_start_threadpool_mock.assert_called_with(
thread_pool_size=7777
) # for the time being, we should have num_threads = num_db_connections
def test_env_priority(mocked_connections, monkeypatch):
"""Env vars should be used over defaults in connect"""
# Use monkeypatch to set environment variable only for this test
monkeypatch.setenv("LOG_LEVEL", "DUMMY_ENV_LOG_LEVEL")
monkeypatch.setenv("LOG_FILE", "DUMMY_ENV_LOG_FILE")
monkeypatch.setenv("DB_PORT", 6969)
monkeypatch.setenv("DB_USER", "DUMMY_ENV_DB_USER")
monkeypatch.setenv("DB_PW", "DUMMY_ENV_DB_PW")
monkeypatch.setenv("DB_HOST", "DUMMY_ENV_DB_HOST")
monkeypatch.setenv("DB_NAME", "DUMMY_ENV_DB_NAME")
monkeypatch.setenv("POOL_SIZE", 7777)
monkeypatch.setenv("POOL_OVERFLOW", 2020)
monkeypatch.setenv("REDIS_HOST", "DUMMY_ENV_REDIS_HOST")
monkeypatch.setenv("REDIS_PORT", 5050)
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect()
core_init_logging_mock.assert_called_with("DUMMY_ENV_LOG_LEVEL",
"DUMMY_ENV_LOG_FILE")
core_init_Connection_mock.assert_called_with(
6969,
"DUMMY_ENV_DB_USER",
"DUMMY_ENV_DB_PW",
"DUMMY_ENV_DB_HOST",
"DUMMY_ENV_DB_NAME",
7777,
2020,
)
core_init_StrictRedis_mock.assert_called_with(host="DUMMY_ENV_REDIS_HOST",
port=5050)
def test_connect_defaults(mocked_connections, monkeypatch):
"""Test connect defaults are used with no params and no env vars"""
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect()
core_init_logging_mock.assert_called_with("error", False)
core_init_Connection_mock.assert_called_with(9000, "analyst", "foo",
"localhost", "flowdb", 5, 1)
core_init_StrictRedis_mock.assert_called_with(host="localhost", port=6379)
def test_log_file_created(tmpdir, monkeypatch):
"""Test that a file handler is created when a path to LOG_FILE is set."""
monkeypatch.setenv("LOG_FILE", str(tmpdir + "log_file"))
connect()
logging.getLogger("flowmachine").error("TEST_MESSAGE")
with open(tmpdir + "log_file") as fin:
log_lines = fin.readline()
assert "TEST_MESSAGE" in log_lines
assert 2 == len(logging.getLogger("flowmachine").handlers)
def main():
# Set up internals and connect to flowdb
flowmachine.connect()
# Set debug mode if required
debug_mode = "true" == os.getenv("DEBUG", "false").lower()
if debug_mode:
logger.info("Enabling asyncio's debugging mode.")
# Run receive loop which receives zmq messages and sends back replies
port = os.getenv("FLOWMACHINE_PORT", 5555)
asyncio.run(recv(port),
debug=debug_mode) # note: asyncio.run() requires Python 3.7+
def fm_conn():
"""
Create a flowmachine Connection object, and connect flowmachine.
Yields
------
flowmachine.core.connection.Connection
"""
fm_conn = make_flowmachine_connection_object()
flowmachine.connect(conn=fm_conn)
yield fm_conn
fm_conn.close()
def main():
# Read config options from environment variables
config = get_server_config()
# Connect to flowdb
flowmachine.connect()
if not config.store_dependencies:
logger.info("Dependency caching is disabled.")
if config.debug_mode:
logger.info("Enabling asyncio's debugging mode.")
# Run receive loop which receives zmq messages and sends back replies
asyncio.run(
recv(config=config),
debug=config.debug_mode,
) # note: asyncio.run() requires Python 3.7+
def flowmachine_connect():
con = flowmachine.connect()
yield con
reset_cache(con)
con.engine.dispose() # Close the connection
Query.redis.flushdb() # Empty the redis
del Query.connection # Ensure we recreate everything at next use
def flowmachine_connect():
con = flowmachine.connect()
yield con
for q in Query.get_stored(): # Remove any cached queries
q.invalidate_db_cache()
con.engine.dispose() # Close the connection
Query.redis.flushdb() # Empty the redis
del Query.connection # Ensure we recreate everything at next use
def fm_conn():
"""
Returns a flowmachine Connection object which
Returns
-------
flowmachine.core.connection.Connection
"""
FLOWDB_HOST = os.getenv("FLOWDB_HOST", "localhost")
FLOWDB_PORT = os.getenv("FLOWDB_PORT", "9000")
conn_str = f"postgresql://*****:*****@{FLOWDB_HOST}:{FLOWDB_PORT}/flowdb"
fm_conn = Connection(conn_str=conn_str)
flowmachine.connect(conn=fm_conn)
yield fm_conn
fm_conn.close()
def test_param_priority(mocked_connections, monkeypatch):
"""Explicit parameters to connect should be respected"""
# Use monkeypatch to set environment variable only for this test
monkeypatch.setenv("LOG_LEVEL", "DUMMY_ENV_LOG_LEVEL")
monkeypatch.setenv("LOG_FILE", "DUMMY_ENV_LOG_FILE")
monkeypatch.setenv("DB_PORT", 7777)
monkeypatch.setenv("DB_USER", "DUMMY_ENV_DB_USER")
monkeypatch.setenv("DB_PW", "DUMMY_ENV_DB_PW")
monkeypatch.setenv("DB_HOST", "DUMMY_ENV_DB_HOST")
monkeypatch.setenv("DB_NAME", "DUMMY_ENV_DB_NAME")
monkeypatch.setenv("POOL_SIZE", 7777)
monkeypatch.setenv("POOL_OVERFLOW", 7777)
monkeypatch.setenv("REDIS_HOST", "DUMMY_ENV_REDIS_HOST")
monkeypatch.setenv("REDIS_PORT", 7777)
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect(
log_level="dummy_log_level",
log_file="dummy_log_file",
db_port=1234,
db_user="******",
db_pw="dummy_db_pw",
db_host="dummy_db_host",
db_name="dummy_db_name",
pool_size=6789,
pool_overflow=1011,
redis_host="dummy_redis_host",
redis_port=1213,
)
core_init_logging_mock.assert_called_with("dummy_log_level",
"dummy_log_file")
core_init_Connection_mock.assert_called_with(
1234,
"dummy_db_user",
"dummy_db_pw",
"dummy_db_host",
"dummy_db_name",
6789,
1011,
)
core_init_StrictRedis_mock.assert_called_with(host="dummy_redis_host",
port=1213)
def test_param_priority(mocked_connections, monkeypatch):
"""Explicit parameters to connect should be respected"""
# Use monkeypatch to set environment variable only for this test
monkeypatch.setenv("LOG_LEVEL", "DUMMY_ENV_LOG_LEVEL")
monkeypatch.setenv("FLOWDB_PORT", "7777")
monkeypatch.setenv("FLOWDB_USER", "DUMMY_ENV_FLOWDB_USER")
monkeypatch.setenv("FLOWDB_PASS", "DUMMY_ENV_FLOWDB_PASS")
monkeypatch.setenv("FLOWDB_HOST", "DUMMY_ENV_FLOWDB_HOST")
monkeypatch.setenv("DB_CONNECTION_POOL_SIZE", "7777")
monkeypatch.setenv("DB_CONNECTION_POOL_OVERFLOW", "7777")
monkeypatch.setenv("REDIS_HOST", "DUMMY_ENV_REDIS_HOST")
monkeypatch.setenv("REDIS_PORT", "7777")
monkeypatch.setenv("REDIS_PASSWORD", "DUMMY_ENV_REDIS_PASSWORD")
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect(
log_level="dummy_log_level",
db_port=1234,
db_user="******",
db_pass="******",
db_host="dummy_db_host",
db_connection_pool_size=6789,
db_connection_pool_overflow=1011,
redis_host="dummy_redis_host",
redis_port=1213,
redis_password="******",
)
core_init_logging_mock.assert_called_with("dummy_log_level")
core_init_Connection_mock.assert_called_with(
port=1234,
user="******",
password="******",
host="dummy_db_host",
database="flowdb",
pool_size=6789,
overflow=1011,
)
core_init_StrictRedis_mock.assert_called_with(
host="dummy_redis_host", port=1213, password="******")
core_init_start_threadpool_mock.assert_called_with(
thread_pool_size=6789
) # for the time being, we should have num_threads = num_db_connections
def fm_conn():
"""
Returns a flowmachine Connection object which
Returns
-------
flowmachine.core.connection.Connection
"""
POSTGRES_USER = os.getenv("POSTGRES_USER", "flowdb")
POSTGRES_PASSWORD = os.getenv("POSTGRES_PASSWORD", "flowflow")
FLOWDB_HOST = os.getenv("FLOWDB_HOST", "localhost")
FLOWDB_PORT = os.getenv("FLOWDB_PORT", "9000")
conn_str = f"postgresql://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{FLOWDB_HOST}:{FLOWDB_PORT}/flowdb"
fm_conn = Connection(conn_str=conn_str)
flowmachine.connect(conn=fm_conn)
yield fm_conn
fm_conn.close()
def test_connect_defaults(mocked_connections, monkeypatch):
"""Test connect defaults are used with no params and no env vars"""
core_init_logging_mock, core_init_Connection_mock, core_init_StrictRedis_mock, core_init_start_threadpool_mock = (
mocked_connections)
connect(db_pass="******", redis_password="******")
core_init_logging_mock.assert_called_with("error")
core_init_Connection_mock.assert_called_with(
port=9000,
user="******",
password="******",
host="localhost",
database="flowdb",
pool_size=5,
overflow=1,
)
core_init_StrictRedis_mock.assert_called_with(host="localhost",
port=6379,
password="******")
core_init_start_threadpool_mock.assert_called_with(
thread_pool_size=5
) # for the time being, we should have num_threads = num_db_connections
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
import flowmachine
import numpy as np
import pandas as pd
import pytest
from flowmachine.core import CustomQuery
from flowmachine.core.subscriber_subsetter import *
from flowmachine.features import daily_location
flowmachine.connect()
# TODO: This is necessary because it is not possible to use a Query object without a connection,
# and no connection is available when the parameterised tests are collected
@pytest.mark.parametrize(
"input_subset, expected_subsetter_type",
[
("all", SubscriberSubsetterForAllSubscribers),
(None, SubscriberSubsetterForAllSubscribers),
("<SINGLE_SUBSCRIBER_ID>", SubscriberSubsetterForExplicitSubset),
(
["<SUBSCRIBER_ID_1>", "<SUBSCRIBER_ID_2>", "<SUBSCRIBER_ID_3>"],
SubscriberSubsetterForExplicitSubset,
),
(
("<SUBSCRIBER_ID_1>", "<SUBSCRIBER_ID_2>"),
SubscriberSubsetterForExplicitSubset,
),
def test_log_level_set(monkeypatch):
"""Test that a log level can be set via param."""
connect(log_level="critical")
assert logging.CRITICAL == logging.getLogger("flowmachine").level
def test_log_level_set_env(monkeypatch):
"""Test that a log level can be set via env."""
monkeypatch.setenv("LOG_LEVEL", "INFO")
connect()
assert logging.INFO == logging.getLogger("flowmachine").level
def test_bad_log_level_goes_to_error(monkeypatch):
"""Test that a bad log level is coerced to ERROR."""
monkeypatch.setenv("LOG_LEVEL", "BAD_LEVEL")
connect()
assert logging.ERROR == logging.getLogger("flowmachine").level
def test_double_connect_warning(monkeypatch):
"""Test that a warning is raised when connecting twice."""
connect()
with pytest.warns(UserWarning):
connect()
assert 1 == len(logging.getLogger("flowmachine").handlers)
def test_connect_passwords_required(args):
"""Test connect raises a valueerror if no password is set for db or redis"""
with pytest.raises(ValueError):
connect(**args)
