redash/tests/factories.py
2019-03-14 10:51:30 +02:00

343 lines
11 KiB
Python

from passlib.apps import custom_app_context as pwd_context
import redash.models
from redash.models import db
from redash.permissions import ACCESS_TYPE_MODIFY
from redash.utils import gen_query_hash, utcnow
from redash.utils.configuration import ConfigurationContainer
class ModelFactory(object):
def __init__(self, model, **kwargs):
self.model = model
self.kwargs = kwargs
def _get_kwargs(self, override_kwargs):
kwargs = self.kwargs.copy()
kwargs.update(override_kwargs)
for key, arg in kwargs.items():
if callable(arg):
kwargs[key] = arg()
return kwargs
def create(self, **override_kwargs):
kwargs = self._get_kwargs(override_kwargs)
obj = self.model(**kwargs)
db.session.add(obj)
db.session.commit()
return obj
class Sequence(object):
def __init__(self, string):
self.sequence = 0
self.string = string
def __call__(self):
self.sequence += 1
return self.string.format(self.sequence)
user_factory = ModelFactory(redash.models.User,
name='John Doe', email=Sequence(u'test{}@example.com'),
password_hash=pwd_context.encrypt('test1234'),
group_ids=[2],
org_id=1)
org_factory = ModelFactory(redash.models.Organization,
name=Sequence("Org {}"),
slug=Sequence("org{}.example.com"),
settings={})
data_source_factory = ModelFactory(redash.models.DataSource,
name=Sequence('Test {}'),
type='pg',
# If we don't use lambda here it will reuse the same options between tests:
options=lambda: ConfigurationContainer.from_json('{"dbname": "test"}'),
org_id=1)
dashboard_factory = ModelFactory(redash.models.Dashboard,
name='test',
user=user_factory.create,
layout='[]',
is_draft=False,
org=1)
api_key_factory = ModelFactory(redash.models.ApiKey,
object=dashboard_factory.create)
query_factory = ModelFactory(redash.models.Query,
name='Query',
description='',
query_text=u'SELECT 1',
user=user_factory.create,
is_archived=False,
is_draft=False,
schedule=None,
data_source=data_source_factory.create,
org_id=1)
query_with_params_factory = ModelFactory(redash.models.Query,
name='New Query with Params',
description='',
query_text='SELECT {{param1}}',
user=user_factory.create,
is_archived=False,
is_draft=False,
schedule={},
data_source=data_source_factory.create,
org_id=1)
access_permission_factory = ModelFactory(redash.models.AccessPermission,
object_id=query_factory.create,
object_type=redash.models.Query.__name__,
access_type=ACCESS_TYPE_MODIFY,
grantor=user_factory.create,
grantee=user_factory.create)
alert_factory = ModelFactory(redash.models.Alert,
name=Sequence('Alert {}'),
query_rel=query_factory.create,
user=user_factory.create,
options={})
query_result_factory = ModelFactory(redash.models.QueryResult,
data='{"columns":{}, "rows":[]}',
runtime=1,
retrieved_at=utcnow,
query_text="SELECT 1",
query_hash=gen_query_hash('SELECT 1'),
data_source=data_source_factory.create,
org_id=1)
visualization_factory = ModelFactory(redash.models.Visualization,
type='CHART',
query_rel=query_factory.create,
name='Chart',
description='',
options='{}')
widget_factory = ModelFactory(redash.models.Widget,
width=1,
options='{}',
dashboard=dashboard_factory.create,
visualization=visualization_factory.create)
destination_factory = ModelFactory(redash.models.NotificationDestination,
org_id=1,
user=user_factory.create,
name=Sequence('Destination {}'),
type='slack',
options=ConfigurationContainer.from_json('{"url": "https://www.slack.com"}'))
alert_subscription_factory = ModelFactory(redash.models.AlertSubscription,
user=user_factory.create,
destination=destination_factory.create,
alert=alert_factory.create)
query_snippet_factory = ModelFactory(redash.models.QuerySnippet,
trigger=Sequence('trigger {}'),
description='description',
snippet='snippet')
class Factory(object):
def __init__(self):
self.org, self.admin_group, self.default_group = redash.models.init_db()
self._data_source = None
self._user = None
@property
def user(self):
if self._user is None:
self._user = self.create_user()
# Test setup creates users, they need to be in the db by the time
# the handler's db transaction starts.
db.session.commit()
return self._user
@property
def data_source(self):
if self._data_source is None:
self._data_source = data_source_factory.create(org=self.org)
db.session.add(redash.models.DataSourceGroup(
group=self.default_group,
data_source=self._data_source))
return self._data_source
def create_org(self, **kwargs):
org = org_factory.create(**kwargs)
self.create_group(org=org, type=redash.models.Group.BUILTIN_GROUP, name="default")
self.create_group(org=org, type=redash.models.Group.BUILTIN_GROUP, name="admin",
permissions=["admin"])
return org
def create_user(self, **kwargs):
args = {
'org': self.org,
'group_ids': [self.default_group.id]
}
if 'org' in kwargs:
args['group_ids'] = [kwargs['org'].default_group.id]
args.update(kwargs)
return user_factory.create(**args)
def create_admin(self, **kwargs):
args = {
'org': self.org,
'group_ids': [self.admin_group.id, self.default_group.id]
}
if 'org' in kwargs:
args['group_ids'] = [kwargs['org'].default_group.id, kwargs['org'].admin_group.id]
args.update(kwargs)
return user_factory.create(**args)
def create_group(self, **kwargs):
args = {
'name': 'Group',
'org': self.org
}
args.update(kwargs)
g = redash.models.Group(**args)
return g
def create_alert(self, **kwargs):
args = {
'user': self.user,
'query_rel': self.create_query()
}
args.update(**kwargs)
return alert_factory.create(**args)
def create_alert_subscription(self, **kwargs):
args = {
'user': self.user,
'alert': self.create_alert()
}
args.update(**kwargs)
return alert_subscription_factory.create(**args)
def create_data_source(self, **kwargs):
group = None
if 'group' in kwargs:
group = kwargs.pop('group')
args = {
'org': self.org
}
args.update(kwargs)
if group and 'org' not in kwargs:
args['org'] = group.org
view_only = args.pop('view_only', False)
data_source = data_source_factory.create(**args)
if group:
db.session.add(redash.models.DataSourceGroup(
group=group,
data_source=data_source,
view_only=view_only))
return data_source
def create_dashboard(self, **kwargs):
args = {
'user': self.user,
'org': self.org
}
args.update(kwargs)
return dashboard_factory.create(**args)
def create_query(self, **kwargs):
args = {
'user': self.user,
'data_source': self.data_source,
'org': self.org
}
args.update(kwargs)
return query_factory.create(**args)
def create_query_with_params(self, **kwargs):
args = {
'user': self.user,
'data_source': self.data_source,
'org': self.org
}
args.update(kwargs)
return query_with_params_factory.create(**args)
def create_access_permission(self, **kwargs):
args = {
'grantor': self.user
}
args.update(kwargs)
return access_permission_factory.create(**args)
def create_query_result(self, **kwargs):
args = {
'data_source': self.data_source,
}
args.update(kwargs)
if 'data_source' in args and 'org' not in args:
args['org'] = args['data_source'].org
return query_result_factory.create(**args)
def create_visualization(self, **kwargs):
args = {
'query_rel': self.create_query()
}
args.update(kwargs)
return visualization_factory.create(**args)
def create_visualization_with_params(self, **kwargs):
args = {
'query_rel': self.create_query_with_params()
}
args.update(kwargs)
return visualization_factory.create(**args)
def create_widget(self, **kwargs):
args = {
'dashboard': self.create_dashboard(),
'visualization': self.create_visualization()
}
args.update(kwargs)
return widget_factory.create(**args)
def create_api_key(self, **kwargs):
args = {
'org': self.org
}
args.update(kwargs)
return api_key_factory.create(**args)
def create_destination(self, **kwargs):
args = {
'org': self.org
}
args.update(kwargs)
return destination_factory.create(**args)
def create_query_snippet(self, **kwargs):
args = {
'user': self.user,
'org': self.org
}
args.update(kwargs)
return query_snippet_factory.create(**args)