import event_model
import xarray
from intake_bluesky.core import (documents_to_xarray,
                                 event_page_to_dataarray_page,
                                 concat_dataarray_pages,
                                 dataarray_page_to_dataset_page)


def no_event_pages(descriptor_uid):
    yield from ()


def event_page_gen(page_size, num_pages):
    """
    Generator event_pages for testing.
    """
    data_keys = ['x', 'y', 'z']
    array_keys = ['seq_num', 'time', 'uid']
    for i in range(num_pages):
        yield {'descriptor': 'DESCRIPTOR',
               **{key: list(range(page_size)) for key in array_keys},
               'data': {key: list(range(page_size)) for key in data_keys},
               'timestamps': {key: list(range(page_size)) for key in data_keys},
               'filled': {key: list(range(page_size)) for key in data_keys}}


def test_no_descriptors():
    run_bundle = event_model.compose_run()
    start_doc = run_bundle.start_doc
    stop_doc = run_bundle.compose_stop()
    documents_to_xarray(
        start_doc=start_doc,
        stop_doc=stop_doc,
        descriptor_docs=[],
        get_event_pages=no_event_pages,
        filler=event_model.Filler({}),
        get_resource=None,
        lookup_resource_for_datum=None,
        get_datum_pages=None)


def test_no_events():
    run_bundle = event_model.compose_run()
    start_doc = run_bundle.start_doc
    desc_bundle = run_bundle.compose_descriptor(
        data_keys={'x': {'source': '...', 'shape': [], 'dtype': 'number'}},
        name='primary')
    descriptor_doc = desc_bundle.descriptor_doc
    stop_doc = run_bundle.compose_stop()
    documents_to_xarray(
        start_doc=start_doc,
        stop_doc=stop_doc,
        descriptor_docs=[descriptor_doc],
        get_event_pages=no_event_pages,
        filler=event_model.Filler({}),
        get_resource=None,
        lookup_resource_for_datum=None,
        get_datum_pages=None)


def test_xarray_helpers():
    event_pages = list(event_page_gen(10, 5))
    dataarray_pages = [event_page_to_dataarray_page(page) for page in event_pages]
    dataarray_page = concat_dataarray_pages(dataarray_pages)
    dataset_page = dataarray_page_to_dataset_page(dataarray_page)
    assert isinstance(dataset_page['data'], xarray.Dataset)
    assert isinstance(dataset_page['timestamps'], xarray.Dataset)
    assert isinstance(dataset_page['filled'], xarray.Dataset)

@@ -26,20 +26,22 @@
  xarray%0A
-from
+import
  intake_
@@ -57,221 +57,15 @@
 ore 
-import (documents_to_xarray,%0A                                 event_page_to_dataarray_page,%0A                                 concat_dataarray_pages,%0A                                 dataarray_page_to_dataset_page)
+as core
 %0A%0A%0Ad
@@ -392,32 +392,51 @@
 key: list(range(
+i*page_size, (i+1)*
 page_size)) for 
@@ -2217,28 +2217,335 @@
 %5B'filled'%5D, xarray.Dataset)%0A
+%0A%0Adef test_interlace_event_page_chunks():%0A    page_gens = %5Bevent_page_gen(10,5) for i in range(3)%5D%0A    interlaced = core.interlace_event_page_chunks(*page_gens, chunk_size=3)%0A%0A    t0 = None%0A    for chunk in interlaced:%0A        t1 = chunk%5B'time'%5D%5B0%5D%0A        if t0:%0A            assert t1 %3E t0%0A        t0 = t1%0A

update to import