From 73927e2b57a28f3c945fd08db762da893be2d274 Mon Sep 17 00:00:00 2001
From: Kay Kurokawa <kay@lbry.io>
Date: Wed, 14 Jun 2017 17:28:37 -0400
Subject: [PATCH] improve tests for GetStream

---
 lbrynet/daemon/Downloader.py                 |   2 -
 tests/unit/lbrynet_daemon/test_Downloader.py | 135 ++++++++++++++++++-
 2 files changed, 130 insertions(+), 7 deletions(-)

diff --git a/lbrynet/daemon/Downloader.py b/lbrynet/daemon/Downloader.py
index ff365f726..03a201d67 100644
--- a/lbrynet/daemon/Downloader.py
+++ b/lbrynet/daemon/Downloader.py
@@ -63,7 +63,6 @@ class GetStream(object):
     def _check_status(self, status):
         stop_condition = (status.num_completed > 0 or
                 status.running_status == ManagedEncryptedFileDownloader.STATUS_STOPPED)
-
         if stop_condition and not self.data_downloading_deferred.called:
             self.data_downloading_deferred.callback(True)
         if self.data_downloading_deferred.called:
@@ -75,7 +74,6 @@ class GetStream(object):
         """
         Check if we've got the first data blob in the stream yet
         """
-
         self.timeout_counter += 1
         if self.timeout_counter >= self.timeout:
             if not self.data_downloading_deferred.called:
diff --git a/tests/unit/lbrynet_daemon/test_Downloader.py b/tests/unit/lbrynet_daemon/test_Downloader.py
index 7e47ac233..0795dc047 100644
--- a/tests/unit/lbrynet_daemon/test_Downloader.py
+++ b/tests/unit/lbrynet_daemon/test_Downloader.py
@@ -2,17 +2,18 @@ import types
 import mock
 import json
 from twisted.trial import unittest
-from twisted.internet import defer
+from twisted.internet import defer, task
 
 from lbryschema.claim import ClaimDict
 
 from lbrynet.core import Session, PaymentRateManager, Wallet
+from lbrynet.core.Error import DownloadTimeoutError
 from lbrynet.daemon import Downloader
 from lbrynet.core.StreamDescriptor import StreamDescriptorIdentifier,StreamMetadata
 from lbrynet.lbry_file.client.EncryptedFileOptions import add_lbry_file_to_sd_identifier
-from lbrynet.core.HashBlob import TempBlob
-from lbrynet.core.BlobManager import TempBlobManager
-from lbrynet.file_manager.EncryptedFileDownloader import ManagedEncryptedFileDownloaderFactory
+
+from lbrynet.file_manager.EncryptedFileStatusReport import EncryptedFileStatusReport
+from lbrynet.file_manager.EncryptedFileDownloader import ManagedEncryptedFileDownloader, ManagedEncryptedFileDownloaderFactory
 from lbrynet.daemon.ExchangeRateManager import ExchangeRateManager
 
 from tests.mocks import BlobAvailabilityTracker as DummyBlobAvailabilityTracker
@@ -20,10 +21,48 @@ from tests.mocks import ExchangeRateManager as DummyExchangeRateManager
 from tests.mocks import BTCLBCFeed, USDBTCFeed
 from tests.mocks import mock_conf_settings
 
+class MocDownloader(object):
+    def __init__(self):
+        self.finish_deferred = defer.Deferred(None)
+        self.stop_called = False
+
+        self.name = 'test'
+        self.num_completed = 0
+        self.num_known = 1
+        self.running_status = ManagedEncryptedFileDownloader.STATUS_RUNNING
+
+    @defer.inlineCallbacks
+    def status(self):
+        out = yield EncryptedFileStatusReport(self.name, self.num_completed, self.num_known, self.running_status)
+        defer.returnValue(out)
+
+    def start(self):
+        return self.finish_deferred
+
+    def stop(self):
+        self.stop_called = True
+        self.finish_deferred.callback(True)
+
+def moc_initialize(self,stream_info):
+    self.sd_hash ="d5169241150022f996fa7cd6a9a1c421937276a3275eb912790bd07ba7aec1fac5fd45431d226b8fb402691e79aeb24b"
+    return None
+
+def moc_download_sd_blob(self):
+    return None
+
+def moc_download(self, sd_blob, name, key_fee):
+    self.pay_key_fee(key_fee, name)
+    self.downloader = MocDownloader()
+    self.downloader.start()
+
+def moc_pay_key_fee(self, key_fee, name):
+    self.pay_key_fee_called = True
+
 class GetStreamTests(unittest.TestCase):
 
     def init_getstream_with_mocs(self):
         mock_conf_settings(self)
+
         sd_identifier = mock.Mock(spec=StreamDescriptorIdentifier)
         session = mock.Mock(spec=Session.Session)
         session.wallet = mock.Mock(spec=Wallet.LBRYumWallet)
@@ -37,8 +76,11 @@ class GetStreamTests(unittest.TestCase):
         data_rate = {'currency':"LBC", 'amount':0, 'address':''}
 
         getstream = Downloader.GetStream(sd_identifier, session,
-            exchange_rate_manager, max_key_fee, timeout=10, data_rate=data_rate)
+            exchange_rate_manager, max_key_fee, timeout=3, data_rate=data_rate)
+        getstream.pay_key_fee_called = False
 
+        self.clock = task.Clock()
+        getstream.checker.clock = self.clock
         return getstream
 
     @defer.inlineCallbacks
@@ -56,3 +98,86 @@ class GetStreamTests(unittest.TestCase):
             yield getstream.start(stream_info,name)
 
 
+    @defer.inlineCallbacks
+    def test_sd_blob_download_timeout(self):
+        """
+        test that if download_sd_blob fails due to timeout,
+        DownloadTimeoutError is raised
+        """
+        def download_sd_blob(self):
+            raise DownloadTimeoutError(self.file_name)
+
+        getstream  = self.init_getstream_with_mocs()
+        getstream._initialize = types.MethodType(moc_initialize, getstream)
+        getstream._download_sd_blob = types.MethodType(download_sd_blob, getstream)
+        getstream._download = types.MethodType(moc_download, getstream)
+        getstream.pay_key_fee = types.MethodType(moc_pay_key_fee, getstream)
+        name='test'
+        stream_info = None
+        with self.assertRaises(DownloadTimeoutError):
+            yield getstream.start(stream_info,name)
+        self.assertFalse(getstream.pay_key_fee_called)
+
+
+    @defer.inlineCallbacks
+    def test_timeout(self):
+        """
+        test that timeout (set to 2 here) exception is raised
+        when download times out while downloading first blob, and key fee is paid
+        """
+        getstream  = self.init_getstream_with_mocs()
+        getstream._initialize = types.MethodType(moc_initialize, getstream)
+        getstream._download_sd_blob = types.MethodType(moc_download_sd_blob, getstream)
+        getstream._download = types.MethodType(moc_download, getstream)
+        getstream.pay_key_fee = types.MethodType(moc_pay_key_fee, getstream)
+        name='test'
+        stream_info = None
+        start = getstream.start(stream_info,name)
+        self.clock.advance(1)
+        self.clock.advance(1)
+        with self.assertRaises(DownloadTimeoutError):
+            yield start
+        self.assertTrue(getstream.downloader.stop_called)
+        self.assertTrue(getstream.pay_key_fee_called)
+
+    @defer.inlineCallbacks
+    def test_finish_one_blob(self):
+        """
+        test that if we have 1 completed blob, start() returns
+        and key fee is paid
+        """
+        getstream  = self.init_getstream_with_mocs()
+        getstream._initialize = types.MethodType(moc_initialize, getstream)
+
+        getstream._download_sd_blob = types.MethodType(moc_download_sd_blob, getstream)
+        getstream._download = types.MethodType(moc_download, getstream)
+        getstream.pay_key_fee = types.MethodType(moc_pay_key_fee, getstream)
+        name='test'
+        stream_info = None
+        start = getstream.start(stream_info,name)
+
+        getstream.downloader.num_completed = 1
+        self.clock.advance(1)
+
+        downloader, f_deferred = yield start
+        self.assertTrue(getstream.pay_key_fee_called)
+
+
+    @defer.inlineCallbacks
+    def test_finish_stopped_downloader(self):
+        """
+        test that if we have a stopped downloader, beforfe a blob is downloaded,
+        start() returns
+        """
+        getstream  = self.init_getstream_with_mocs()
+        getstream._initialize = types.MethodType(moc_initialize, getstream)
+        getstream._download_sd_blob = types.MethodType(moc_download_sd_blob, getstream)
+        getstream._download = types.MethodType(moc_download, getstream)
+        name='test'
+        stream_info = None
+        start = getstream.start(stream_info,name)
+
+        getstream.downloader.running_status = ManagedEncryptedFileDownloader.STATUS_STOPPED
+        self.clock.advance(1)
+        downloader, f_deferred = yield start
+