From ea3c21220c7d3cffd4964f99df43bee288c45831 Mon Sep 17 00:00:00 2001
From: Andrew Gorcester <gorcester@google.com>
Date: Tue, 10 Sep 2019 10:47:12 -0700
Subject: [PATCH] Update NDB samples to consolidate tags

---
 .../{middleware.py => django_middleware.py}   | 16 ++-----------
 ...ware_test.py => django_middleware_test.py} | 16 +++----------
 datastore/cloud-ndb/flask_app.py              | 24 ++++++++++++-------
 datastore/cloud-ndb/flask_app_test.py         | 13 ++++++++++
 datastore/cloud-ndb/requirements.txt          |  2 ++
 5 files changed, 35 insertions(+), 36 deletions(-)
 rename datastore/cloud-ndb/{middleware.py => django_middleware.py} (72%)
 rename datastore/cloud-ndb/{middleware_test.py => django_middleware_test.py} (66%)

diff --git a/datastore/cloud-ndb/middleware.py b/datastore/cloud-ndb/django_middleware.py
similarity index 72%
rename from datastore/cloud-ndb/middleware.py
rename to datastore/cloud-ndb/django_middleware.py
index 8cec79b1691..7506cc33334 100644
--- a/datastore/cloud-ndb/middleware.py
+++ b/datastore/cloud-ndb/django_middleware.py
@@ -12,24 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# [START ndb_wsgi_middleware]
 # [START ndb_django_middleware]
 from google.cloud import ndb
 
 
-# [END ndb_django_middleware]
-def ndb_wsgi_middleware(wsgi_app):
-    client = ndb.Client()
-
-    def middleware(environ, start_response):
-        with client.context():
-            return wsgi_app(environ, start_response)
-
-    return middleware
-# [END ndb_wsgi_middleware]
-
-
-# [START ndb_django_middleware]
+# Once this middleware is activated in Django settings, NDB calls inside Django
+# views will be executed in context, with a separate context for each request.
 def ndb_django_middleware(get_response):
     client = ndb.Client()
 
diff --git a/datastore/cloud-ndb/middleware_test.py b/datastore/cloud-ndb/django_middleware_test.py
similarity index 66%
rename from datastore/cloud-ndb/middleware_test.py
rename to datastore/cloud-ndb/django_middleware_test.py
index c6a33720e9a..22c6c65f0e8 100644
--- a/datastore/cloud-ndb/middleware_test.py
+++ b/datastore/cloud-ndb/django_middleware_test.py
@@ -14,18 +14,7 @@
 
 from google.cloud import ndb
 
-import middleware
-
-
-def test_ndb_wsgi_middleware():
-    def fake_wsgi_app(environ, start_response):
-        # Validate that a context is live. This will throw
-        # ndb.exceptions.ContextError if no context is set up.
-        ndb.context.get_context()
-
-    wrapped_function = middleware.ndb_wsgi_middleware(fake_wsgi_app)
-
-    wrapped_function(None, None)
+import django_middleware
 
 
 def test_ndb_django_middleware():
@@ -34,6 +23,7 @@ def fake_get_response(request):
         # ndb.exceptions.ContextError if no context is set up.
         ndb.context.get_context()
 
-    wrapped_function = middleware.ndb_django_middleware(fake_get_response)
+    wrapped_function = django_middleware.ndb_django_middleware(
+        fake_get_response)
 
     wrapped_function(None)
diff --git a/datastore/cloud-ndb/flask_app.py b/datastore/cloud-ndb/flask_app.py
index 2bb7fdd8942..1d1153068cc 100644
--- a/datastore/cloud-ndb/flask_app.py
+++ b/datastore/cloud-ndb/flask_app.py
@@ -12,27 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from middleware import ndb_wsgi_middleware
-
-# [START ndb_flask_middleware_usage]
+# [START ndb_flask]
 from flask import Flask
-# [END ndb_flask_middleware_usage]
+
 from google.cloud import ndb
-# [START ndb_flask_middleware_usage]
+
+
+client = ndb.Client()
+
+
+def ndb_wsgi_middleware(wsgi_app):
+    def middleware(environ, start_response):
+        with client.context():
+            return wsgi_app(environ, start_response)
+
+    return middleware
+
 
 app = Flask(__name__)
 app.wsgi_app = ndb_wsgi_middleware(app.wsgi_app)  # Wrap the app in middleware.
-# [END ndb_flask_middleware_usage]
 
 
 class Book(ndb.Model):
     title = ndb.StringProperty()
 
 
-client = ndb.Client()
-
-
 @app.route('/')
 def list_books():
     books = Book.query()
     return str([book.to_dict() for book in books])
+# [END ndb_flask]
diff --git a/datastore/cloud-ndb/flask_app_test.py b/datastore/cloud-ndb/flask_app_test.py
index 4226f8e554e..8296a1ef09f 100644
--- a/datastore/cloud-ndb/flask_app_test.py
+++ b/datastore/cloud-ndb/flask_app_test.py
@@ -16,6 +16,8 @@
 
 import pytest
 
+from google.cloud import ndb
+
 import flask_app
 
 
@@ -40,3 +42,14 @@ def test_index(test_book):
     r = client.get('/')
     assert r.status_code == 200
     assert test_book.title in r.data.decode('utf-8')
+
+
+def test_ndb_wsgi_middleware():
+    def fake_wsgi_app(environ, start_response):
+        # Validate that a context is live. This will throw
+        # ndb.exceptions.ContextError if no context is set up.
+        ndb.context.get_context()
+
+    wrapped_function = flask_app.ndb_wsgi_middleware(fake_wsgi_app)
+
+    wrapped_function(None, None)
diff --git a/datastore/cloud-ndb/requirements.txt b/datastore/cloud-ndb/requirements.txt
index b92fbc0bd38..fa41f8a3e5b 100644
--- a/datastore/cloud-ndb/requirements.txt
+++ b/datastore/cloud-ndb/requirements.txt
@@ -1,2 +1,4 @@
+#[START ndb_version]
 google-cloud-ndb==0.0.1
+#[END ndb_version]
 flask==1.0.2
\ No newline at end of file