From 94a70e9ad1c6bacb5a5ed547c1689c4eae321677 Mon Sep 17 00:00:00 2001
From: Neil Greenwood <nbartongreenwood@alertlogic.com>
Date: Tue, 4 May 2021 17:45:23 +0100
Subject: [PATCH] `nextToken` value in `logs:describeLogStreams` response
 (#3896)

* `nextToken` value in `logs:describeLogStreams` response

Modified the pagination for FilterLogEvents to more closely follow
the real AWS behaviour.

* Make assertions work in py2 and py3.
---
 moto/logs/models.py          | 34 ++++++++++++++++-----
 tests/test_logs/test_logs.py | 57 ++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 7 deletions(-)

diff --git a/moto/logs/models.py b/moto/logs/models.py
index 57638653..030f20b6 100644
--- a/moto/logs/models.py
+++ b/moto/logs/models.py
@@ -289,14 +289,34 @@ class LogGroup:
                 else item[1].get("lastEventTimestamp", 0)
             )
 
-        if next_token is None:
-            next_token = 0
-
         log_streams = sorted(log_streams, key=sorter, reverse=descending)
-        new_token = next_token + limit
-        log_streams_page = [x[1] for x in log_streams[next_token:new_token]]
-        if new_token >= len(log_streams):
-            new_token = None
+        first_index = 0
+        if next_token:
+            try:
+                group, stream = next_token.split("@")
+                if group != log_group_name:
+                    raise ValueError()
+                first_index = (
+                    next(
+                        index
+                        for (index, e) in enumerate(log_streams)
+                        if e[1]["logStreamName"] == stream
+                    )
+                    + 1
+                )
+            except (ValueError, StopIteration):
+                first_index = 0
+                log_streams = []
+
+        last_index = first_index + limit
+        if last_index > len(log_streams):
+            last_index = len(log_streams)
+        log_streams_page = [x[1] for x in log_streams[first_index:last_index]]
+        new_token = None
+        if log_streams_page and last_index < len(log_streams):
+            new_token = "{}@{}".format(
+                log_group_name, log_streams_page[-1]["logStreamName"]
+            )
 
         return log_streams_page, new_token
 
diff --git a/tests/test_logs/test_logs.py b/tests/test_logs/test_logs.py
index 7bad73c5..d61b457b 100644
--- a/tests/test_logs/test_logs.py
+++ b/tests/test_logs/test_logs.py
@@ -554,3 +554,60 @@ def test_describe_log_groups_paging():
     resp = client.describe_log_groups(nextToken="invalid-token")
     resp["logGroups"].should.have.length_of(0)
     resp.should_not.have.key("nextToken")
+
+
+@mock_logs
+def test_describe_log_streams_paging():
+    client = boto3.client("logs", "us-east-1")
+
+    log_group_name = "/aws/codebuild/lowercase-dev"
+    stream_names = [
+        "job/214/stage/unit_tests/foo",
+        "job/215/stage/unit_tests/spam",
+        "job/215/stage/e2e_tests/eggs",
+        "job/216/stage/unit_tests/eggs",
+    ]
+
+    client.create_log_group(logGroupName=log_group_name)
+    for name in stream_names:
+        client.create_log_stream(logGroupName=log_group_name, logStreamName=name)
+
+    resp = client.describe_log_streams(logGroupName=log_group_name)
+    resp["logStreams"].should.have.length_of(4)
+    resp["logStreams"][0]["arn"].should.contain(log_group_name)
+    resp.should_not.have.key("nextToken")
+
+    resp = client.describe_log_streams(logGroupName=log_group_name, limit=2)
+    resp["logStreams"].should.have.length_of(2)
+    resp["logStreams"][0]["arn"].should.contain(log_group_name)
+    resp["nextToken"].should.equal(
+        u"{}@{}".format(log_group_name, resp["logStreams"][1]["logStreamName"])
+    )
+
+    resp = client.describe_log_streams(
+        logGroupName=log_group_name, nextToken=resp["nextToken"], limit=1
+    )
+    resp["logStreams"].should.have.length_of(1)
+    resp["logStreams"][0]["arn"].should.contain(log_group_name)
+    resp["nextToken"].should.equal(
+        u"{}@{}".format(log_group_name, resp["logStreams"][0]["logStreamName"])
+    )
+
+    resp = client.describe_log_streams(
+        logGroupName=log_group_name, nextToken=resp["nextToken"]
+    )
+    resp["logStreams"].should.have.length_of(1)
+    resp["logStreams"][0]["arn"].should.contain(log_group_name)
+    resp.should_not.have.key("nextToken")
+
+    resp = client.describe_log_streams(
+        logGroupName=log_group_name, nextToken="invalid-token"
+    )
+    resp["logStreams"].should.have.length_of(0)
+    resp.should_not.have.key("nextToken")
+
+    resp = client.describe_log_streams(
+        logGroupName=log_group_name, nextToken="invalid@token"
+    )
+    resp["logStreams"].should.have.length_of(0)
+    resp.should_not.have.key("nextToken")