factories.py

import factory
import psycopg2

from funkwhale_network import serializers


class DBFactory(factory.Factory):
    class Meta:
        model = dict

    @classmethod
    def set_db(cls, conn, table):
        cls._conn = conn
        cls._table = table

    @classmethod
    async def c(cls, *args, **kwargs):
        o = cls(*args, **kwargs)
        await cls.pre_create(o)
        return await cls.persist_to_db(o)

    @classmethod
    async def pre_create(cls, o):
        pass

    @classmethod
    async def persist_to_db(cls, o):
        fields, values = [], []
        for field, value in o.items():
            fields.append(field)
            values.append(value)

        sql = "INSERT INTO {} ({}) VALUES ({}) RETURNING *".format(
            cls._table, ", ".join(fields), ", ".join(["%s" for _ in values])
        )
        async with cls._conn.cursor(
            cursor_factory=psycopg2.extras.RealDictCursor
        ) as cursor:
            await cursor.execute(sql, values)
            return await cursor.fetchone()


class DomainFactory(DBFactory):
    name = factory.Faker("domain_name")
    node_name = factory.Faker("paragraph")
    blocked = False


class CheckFactory(DBFactory):
    time = "NOW()"
    up = factory.Faker("boolean")
    domain = factory.Faker("domain_name")
    open_registrations = factory.Faker("boolean")
    private = factory.Faker("boolean")
    federation_enabled = factory.Faker("boolean")
    anonymous_can_listen = factory.Faker("boolean")
    usage_users_total = factory.Faker("random_int")
    usage_users_active_half_year = factory.Faker("random_int")
    usage_users_active_month = factory.Faker("random_int")
    usage_listenings_total = factory.Faker("random_int")
    library_tracks_total = factory.Faker("random_int")
    library_albums_total = factory.Faker("random_int")
    library_artists_total = factory.Faker("random_int")
    library_music_hours = factory.Faker("random_int")
    software_name = "funkwhale"
    software_version_major = factory.Faker("random_digit")
    software_version_minor = factory.Faker("random_digit")
    software_version_patch = factory.Faker("random_digit")
    software_prerelease = None
    software_build = None

    @classmethod
    async def pre_create(cls, o):
        await serializers.create_domain(cls._conn, {"name": o["domain"]})


ALL = [("Check", CheckFactory, "checks"), ("Domain", DomainFactory, "domains")]