diff --git a/README.md b/README.md index 90fffe7..b03d9fc 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ import ( ) func main() { - for tweet := range twitterscraper.GetTweets(context.Background(), "Twitter", 5) { + for tweet := range twitterscraper.GetTweets(context.Background(), "Twitter", 50) { if tweet.Error != nil { panic(tweet.Error) } @@ -37,7 +37,7 @@ func main() { } ``` -It appears you can ask for up to 5 pages of tweets reliably. +It appears you can ask for up to 50 tweets. ### Search tweets by query standard operators diff --git a/tweets.go b/tweets.go index 1447279..fc1a9f4 100644 --- a/tweets.go +++ b/tweets.go @@ -45,12 +45,13 @@ type Result struct { } // GetTweets returns channel with tweets for a given user. -func GetTweets(ctx context.Context, user string, pages int) <-chan *Result { +func GetTweets(ctx context.Context, user string, maxTweetsNbr int) <-chan *Result { channel := make(chan *Result) go func(user string) { defer close(channel) var lastTweetID string - for pages > 0 { + tweetsNbr := 0 + for tweetsNbr < maxTweetsNbr { select { case <-ctx.Done(): channel <- &Result{Error: ctx.Err()} @@ -63,6 +64,11 @@ func GetTweets(ctx context.Context, user string, pages int) <-chan *Result { channel <- &Result{Error: err} return } + + if len(tweets) == 0 { + break + } + for _, tweet := range tweets { select { case <-ctx.Done(): @@ -71,10 +77,12 @@ func GetTweets(ctx context.Context, user string, pages int) <-chan *Result { default: } - lastTweetID = tweet.ID - channel <- &Result{Tweet: *tweet} + if tweetsNbr < maxTweetsNbr { + lastTweetID = tweet.ID + channel <- &Result{Tweet: *tweet} + } + tweetsNbr++ } - pages-- } }(user) return channel diff --git a/tweets_test.go b/tweets_test.go index 90af459..fd2ed51 100644 --- a/tweets_test.go +++ b/tweets_test.go @@ -7,7 +7,8 @@ import ( func TestGetTweets(t *testing.T) { count := 0 - for tweet := range GetTweets(context.Background(), "nomadic_ua", 2) { + maxTweetsNbr := 50 + for tweet := range GetTweets(context.Background(), "Twitter", maxTweetsNbr) { if tweet.Error != nil { t.Error(tweet.Error) } else { @@ -32,7 +33,7 @@ func TestGetTweets(t *testing.T) { } } } - if count == 0 { - t.Error("Expected tweets count is greater than zero") + if count != maxTweetsNbr { + t.Errorf("Expected tweets count=%v, got: %v", maxTweetsNbr, count) } }