diff --git a/fuelclient/cli/actions/snapshot.py b/fuelclient/cli/actions/snapshot.py index f2e77e1..993458c 100644 --- a/fuelclient/cli/actions/snapshot.py +++ b/fuelclient/cli/actions/snapshot.py @@ -19,6 +19,7 @@ import yaml from fuelclient.cli.actions.base import Action import fuelclient.cli.arguments as Args from fuelclient.cli.formatting import download_snapshot_with_progress_bar +from fuelclient.client import APIClient from fuelclient.objects.task import SnapshotTask @@ -63,6 +64,7 @@ class SnapshotAction(Action): snapshot_task.wait() download_snapshot_with_progress_bar( snapshot_task.connection.root + snapshot_task.data["message"], + auth_token=APIClient.auth_token, directory=params.dir ) diff --git a/fuelclient/cli/formatting.py b/fuelclient/cli/formatting.py index dbec002..300165f 100644 --- a/fuelclient/cli/formatting.py +++ b/fuelclient/cli/formatting.py @@ -103,7 +103,8 @@ def get_bar_for_progress(full_width, progress): ) -def download_snapshot_with_progress_bar(url, directory=os.path.curdir): +def download_snapshot_with_progress_bar( + url, auth_token, directory=os.path.curdir): """downloads file from specific 'url' with progress bar and save it to some 'directory'. """ @@ -113,7 +114,8 @@ def download_snapshot_with_progress_bar(url, directory=os.path.curdir): os.path.abspath(directory), url.split('/')[-1] ) - download_handle = urllib2.urlopen(url) + request = urllib2.Request(url, headers={'x-auth-token': auth_token}) + download_handle = urllib2.urlopen(request) with open(file_name, 'wb') as file_handle: meta = download_handle.info() file_size = int(meta.getheaders("Content-Length")[0]) diff --git a/fuelclient/tests/unit/v1/test_snapshot.py b/fuelclient/tests/unit/v1/test_snapshot.py index 2316be5..6028d94 100644 --- a/fuelclient/tests/unit/v1/test_snapshot.py +++ b/fuelclient/tests/unit/v1/test_snapshot.py @@ -12,7 +12,7 @@ # License for the specific language governing permissions and limitations # under the License. - +import mock from mock import call from mock import patch @@ -30,11 +30,13 @@ class TestSnapshot(base.UnitTestCase): self.execute(['fuel', 'snapshot', '--conf']) self.assertEqual(mstdout.write.call_args_list, [call('key: value\n')]) + @patch('fuelclient.cli.actions.snapshot.APIClient', + mock.Mock(auth_token='token123')) @patch('fuelclient.cli.actions.snapshot.SnapshotTask.start_snapshot_task') @patch('fuelclient.cli.actions.snapshot.' 'download_snapshot_with_progress_bar') @patch('sys.stdin') - def test_snapshot_with_provided_conf(self, mstdin, mbar, mstart): + def test_snapshot_with_provided_conf(self, mstdin, mdownload, mstart): conf = 'key: value\n' mstdin.isatty.return_value = False @@ -45,14 +47,25 @@ class TestSnapshot(base.UnitTestCase): mstart.assert_called_once_with({'key': 'value'}) self.assertEqual(mstdin.read.call_count, 1) + mdownload.assert_called_once_with( + mock.ANY, + auth_token='token123', + directory='.') + + @patch('fuelclient.cli.actions.snapshot.APIClient', + mock.Mock(auth_token='token123')) @patch('fuelclient.cli.actions.snapshot.SnapshotTask.start_snapshot_task') @patch('fuelclient.cli.actions.snapshot.' 'download_snapshot_with_progress_bar') @patch('sys.stdin') - def test_snapshot_without_conf(self, mstdin, mbar, mstart): + def test_snapshot_without_conf(self, mstdin, mdownload, mstart): mstdin.isatty.return_value = True self.execute(['fuel', 'snapshot']) mstart.assert_called_once_with({}) + mdownload.assert_called_once_with( + mock.ANY, + auth_token='token123', + directory='.')