package com.novoda.downloadmanager;

import static com.novoda.downloadmanager.DownloadBatchStatus.Status.DELETION;
import static com.novoda.downloadmanager.DownloadBatchStatus.Status.DOWNLOADED;

class NotificationDispatcher {

    private final Object waitForDownloadService;
    private final NotificationCreator<DownloadBatchStatus> notificationCreator;

    private DownloadService downloadService;

    NotificationDispatcher(Object waitForDownloadService, NotificationCreator<DownloadBatchStatus> notificationCreator) {
        this.waitForDownloadService = waitForDownloadService;
        this.notificationCreator = notificationCreator;
    }

    void updateNotification(DownloadBatchStatus downloadBatchStatus) {
        WaitForDownloadService.<Void>waitFor(downloadService, waitForDownloadService)
                .thenPerform(executeUpdateNotification(downloadBatchStatus));
    }

    private WaitForDownloadService.ThenPerform.Action<Void> executeUpdateNotification(DownloadBatchStatus downloadBatchStatus) {
        return () -> {
            NotificationInformation notificationInformation = notificationCreator.createNotification(downloadBatchStatus);

            if (downloadBatchStatus.status() == DOWNLOADED || downloadBatchStatus.status() == DELETION) {
                downloadService.stackNotification(notificationInformation);
            } else {
                downloadService.updateNotification(notificationInformation);
            }

            return null;
        };
    }

    void setDownloadService(DownloadService downloadService) {
        this.downloadService = downloadService;
    }
}

@@ -30,16 +30,65 @@
 nager;%0A%0A
+import android.support.annotation.WorkerThread;%0A%0A
 import s
@@ -685,32 +685,50 @@
 Creator;%0A    %7D%0A%0A
+    @WorkerThread%0A
     void updateN

Add worker thread annotation to update notification.