kraken/backend/tests/utils/utils.py

"""provide fake data obj"""
from datetime import datetime, timedelta
import itertools
import random

from faker import Faker
from faker.providers import DynamicProvider
from sqlalchemy.orm import Session

from db.models.subjects import Subject
from db.models.trains import Train
from db.models.hand_peripherals import HandPeripheral
from db.models.daily_stats import DailyStats
from db.repository import subjects as db_rep_sub
from schemas.subjects import SubjectCreate


class FakerManager:
    """init fake obj"""

    def __init__(self, lang="zh-cn"):
        self.fake = Faker(lang)
        self.load_provider()

    def load_provider(self):
        self.fake.add_provider(self.get_gender_provider())
        self.fake.add_provider(self.get_parts_provider())
        self.fake.add_provider(self.get_train_rank_provider())

    @staticmethod
    def get_gender_provider():
        return DynamicProvider(provider_name="subject_gender",
                               elements=["男", "女"])

    @staticmethod
    def get_parts_provider():
        parts_list = []
        for num in range(1, 5):
            parts_list.extend(
                list(itertools.combinations(["左手", "右手", "左腿", "右腿"], num)))
        parts_provider = DynamicProvider(provider_name="rehabilitation_parts",
                                         elements=parts_list)
        return parts_provider

    @staticmethod
    def get_train_rank_provider():
        return DynamicProvider(provider_name="train_rank",
                               elements=["简单", "中等", "困难"])



def get_random_position(subject):
    return random.choice(subject.rehabilitation_parts)


def generate_delay_datetime(delay_years: int):
    today = datetime.today()
    delay_time = timedelta(days=365*delay_years)
    delay_datetime = today + delay_time
    return delay_datetime.strftime("%Y-%m-%d")


def clear_db_table(db: Session):
    db.query(Subject).delete()
    db.query(Train).delete()
    db.query(HandPeripheral).delete()
    db.query(DailyStats).delete()
    db.commit()


fake = FakerManager().fake