def test_plot_with_sppeople(create_pop, do_show=False, do_save=False):
    """
    Test plotting method works on synthpops.people.People object.

    Notes:
        With this pop type, you will need to supply more information to
        tell the method where to look for expected data.
    """
    sp.logger.info(
        "Test that the age comparison plotting method works on sp.people.People and plotting styles can be easily updated."
    )
    pop = create_pop
    people = pop.to_people()
    kwargs = sc.objdict(sc.mergedicts(pars, pop.loc_pars))
    kwargs.datadir = sp.settings.datadir
    kwargs.figname = f"test_ages_{kwargs.location}_sppeople"
    kwargs.do_show = do_show
    kwargs.do_save = do_save

    # modify some plotting styles
    kwargs.color_1 = '#9966cc'
    kwargs.color_2 = 'indigo'
    kwargs.markersize = 4.5
    fig, ax = sp.plot_ages(people, **kwargs)
    # fig, ax = sp.plot_ages(people)  # to plot without extra information

    assert isinstance(fig, mplt.figure.Figure), 'Check failed.'
    print('Check passed. Figure made.')

    return fig, ax, people
def test_plot_ages(do_show=False, do_save=False):
    """
    Test that the age comparison plotting method in sp.Pop class works.

    Note:
        With any popdict, you will need to supply more information to
        tell the method where to look for expected data.
    """
    sp.logger.info(
        "Test that the age comparison plotting method with sp.Pop object.")
    pop = sp.Pop(**pars)
    kwargs = sc.objdict(sc.mergedicts(pars, pop.loc_pars))
    kwargs.figname = f"test_pop_ages_{kwargs.location}_pop"
    kwargs.do_show = do_show
    kwargs.do_save = do_save
    fig, ax = pop.plot_ages(**kwargs)
    # fig, ax = pop.plot_ages()  # to plot without extra information

    assert isinstance(fig, mplt.figure.Figure), 'Check 1 failed.'
    print('Check passed. Figure 1 made.')

    popdict = pop.to_dict()
    kwargs.datadir = sp.datadir  # extra information required
    kwargs.figname = f"test_popdict_ages_{kwargs.location}_popdict"
    kwargs.do_show = False
    fig2, ax2 = sp.plot_ages(popdict, **kwargs)
    # fig2, ax2 = sp.plot_ages(popdict)  # to plot without extra information
    if not kwargs.do_show:
        plt.close()
    assert isinstance(fig, mplt.figure.Figure), 'Check 2 failed.'
    print('Check passed. Figure 2 made.')
    return fig, ax, pop
def test_catch_pop_type_errors():
    """
    Test that synthpops.plotting methods raise error when pop type is not in
    sp.Pop, dict, or cv.people.People.
    """
    sp.logger.info("Catch NotImplementedError when pop type is invalid.")
    pop = 'not a pop object'

    with pytest.raises(NotImplementedError):
        sp.plot_ages(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_household_sizes(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_ltcf_resident_sizes(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_enrollment_rates_by_age(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_employment_rates_by_age(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_school_sizes(pop)
    with pytest.raises(NotImplementedError):
        sp.plot_workplace_sizes(pop)
def test_plot_with_cvpeople(create_pop, do_show=False, do_save=False):
    """
    Test plotting method works on covasim.people.People object.

    Notes:
        With this pop type, you will need to supply more information to
        tell the method where to look for expected data.
    """
    sp.logger.info(
        "Test that the age comparison plotting method works on cv.people.People and plotting styles can be easily updated."
    )
    pop = create_pop
    popdict = pop.to_dict()
    cvpopdict = cv.make_synthpop(population=popdict,
                                 community_contacts=2)  # array based

    # Actually create the people
    people_pars = dict(
        pop_size=pars.n,
        beta_layer={k: 1.0
                    for k in 'hswcl'
                    },  # Since this is used to define hat layers exist
        beta=
        1.0,  # TODO: this is required for plotting (people.plot()), but shouldn't be (in covasim)
    )
    people = cv.People(people_pars,
                       strict=False,
                       uid=cvpopdict['uid'],
                       age=cvpopdict['age'],
                       sex=cvpopdict['sex'])
    kwargs = sc.objdict(sc.mergedicts(pars, pop.loc_pars))
    kwargs.datadir = sp.settings.datadir
    kwargs.figname = f"test_ages_{kwargs.location}_cvpeople"
    kwargs.do_show = do_show
    kwargs.do_save = do_save

    # modify some plotting styles
    kwargs.color_1 = '#9966cc'
    kwargs.color_2 = 'indigo'
    kwargs.markersize = 4.5
    fig, ax = sp.plot_ages(people, **kwargs)
    # fig, ax = sp.plot_ages(people)  # to plot without extra information

    assert isinstance(fig, mplt.figure.Figure), 'Check failed.'
    print('Check passed. Figure made.')

    return fig, ax, people