<s> import json <EOL> import unittest <EOL> from boto . s3 . key import Key <EOL> from mock import create_autospec , Mock , call <EOL> from arbalest . s3 import Bucket <EOL> from arbalest . pipeline import S3SortedDataSources <EOL> def mock_key ( name ) : <EOL> return Key ( Mock ( ) , name ) <EOL> class S3SortedDataSourcesShould ( unittest . TestCase ) : <EOL> def setUp ( self ) : <EOL> parents = [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] <EOL> first_children = [ '<STR_LIT>' , <EOL> '<STR_LIT>' ] <EOL> second_children = [ '<STR_LIT>' ] <EOL> self . bucket = create_autospec ( Bucket ) <EOL> self . bucket . list = Mock ( <EOL> side_effect = [ [ mock_key ( key ) for key in parents ] , <EOL> [ mock_key ( key ) for key in first_children ] , <EOL> [ mock_key ( key ) for key in second_children ] ] ) <EOL> def test_have_source_journal_key ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket ) <EOL> self . assertEqual ( '<STR_LIT>' , <EOL> source . source_journal_key ) <EOL> def test_get_all_dates_as_sources ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket ) <EOL> self . assertEqual ( [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] , <EOL> list ( source . get ( ) ) ) <EOL> self . bucket . list . assert_called_once_with ( source . source + '<STR_LIT:/>' , '<STR_LIT:/>' ) <EOL> def test_get_all_dates_as_sources_with_empty_dates ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket , <EOL> '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . assertEqual ( [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] , <EOL> list ( source . get ( ) ) ) <EOL> self . bucket . list . assert_called_once_with ( source . source + '<STR_LIT:/>' , '<STR_LIT:/>' ) <EOL> def test_get_all_dates_including_and_after_start_date_as_sources ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket , <EOL> '<STR_LIT>' ) <EOL> self . assertEqual ( [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] , <EOL> list ( source . get ( ) ) ) <EOL> self . bucket . list . assert_called_once_with ( source . source + '<STR_LIT:/>' , '<STR_LIT:/>' ) <EOL> def test_get_all_dates_including_and_before_end_date_as_sources ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket , <EOL> end = '<STR_LIT>' ) <EOL> self . assertEqual ( [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] , <EOL> list ( source . get ( ) ) ) <EOL> self . bucket . list . assert_called_once_with ( source . source + '<STR_LIT:/>' , '<STR_LIT:/>' ) <EOL> def test_get_all_dates_including_and_between_start_and_end_date_as_sources ( <EOL> self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket , <EOL> start = '<STR_LIT>' , <EOL> end = '<STR_LIT>' ) <EOL> self . assertEqual ( [ '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' ] , <EOL> list ( source . get ( ) ) ) <EOL> self . bucket . list . assert_has_calls ( <EOL> [ call ( source . source + '<STR_LIT:/>' , '<STR_LIT:/>' ) , <EOL> call ( '<STR_LIT>' , '<STR_LIT:/>' ) , <EOL> call ( '<STR_LIT>' , '<STR_LIT:/>' ) , <EOL> call ( '<STR_LIT>' , '<STR_LIT:/>' ) ] ) <EOL> def test_committed ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket ) <EOL> key = mock_key ( source . source_journal_key ) <EOL> key . exists = Mock ( return_value = True ) <EOL> source . bucket . get = Mock ( return_value = key ) <EOL> self . assertEqual ( True , source . committed ( ) . exists ( ) ) <EOL> def test_commit ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket ) <EOL> committed_point = '<STR_LIT>' <EOL> source . commit ( committed_point ) <EOL> self . bucket . save . assert_called_once_with ( source . source_journal_key , <EOL> json . dumps ( { <EOL> '<STR_LIT>' : committed_point } ) ) <EOL> def test_rollback ( self ) : <EOL> source = S3SortedDataSources ( '<STR_LIT>' , '<STR_LIT>' , self . bucket ) <EOL> source . rollback ( ) <EOL> self . bucket . delete . assert_called_once_with ( source . source_journal_key ) </s>
