

import pytest

from osp_pipeline.core.services import config

from osp_pipeline.corpus.jobs.list_v2_paths \
    import main as list_v2_paths

from osp_pipeline.corpus.jobs.load_v2_warcs \
    import main as load_v2_warcs

from tests.utils import read_yaml, data_to_text


cases = read_yaml(__file__, 'test_load_v2_warcs.yml')


@pytest.fixture(scope='module')
def warcs(sc, spark, results, sync_bucket):
    """Run jobs, provide result.
    """
    sync_bucket(config['buckets']['v2_corpus'])

    list_v2_paths(sc, spark)

    load_v2_warcs(sc, spark)

    return spark.read.parquet(config['results']['v2_corpus_warcs'])


@pytest.mark.parametrize('uuid,fields', cases.items())
@pytest.mark.skip(reason="problems with fake-s3.")
def test_load_v2_warcs(warcs, uuid, fields):
    """Extract metadata + text from WARCs.
    """
    row = warcs.filter(warcs.id == uuid).head()

    if 'url' in fields:
        assert row.url == fields['url']

    if 'source_url' in fields:
        assert row.source_url == fields['source_url']

    if 'source_anchor' in fields:
        assert row.source_anchor == fields['source_anchor']

    if 'retrieved' in fields:
        assert row.retrieved == fields['retrieved']

    if 'spider_run_id' in fields:
        assert row.spider_run_id == fields['spider_run_id']

    if 'depth' in fields:
        assert row.depth == fields['depth']

    if 'hops_from_seed' in fields:
        assert row.hops_from_seed == fields['hops_from_seed']

    if 'status' in fields:
        assert row.status == fields['status']

    if 'text' in fields:
        assert fields['text'] in data_to_text(row.data)
