From f6979f5d099c9e54695df4c86614595eb6f50041 Mon Sep 17 00:00:00 2001
From: JJ Geewax <jj@geewax.org>
Date: Tue, 25 Feb 2014 10:39:24 -0500
Subject: [PATCH] Fix #50 - Missing time import.

Also added tests and made a minor modification to Key.to_protobuf().
---
 gcloud/datastore/helpers.py      |  1 +
 gcloud/datastore/key.py          | 20 +++++++------
 gcloud/datastore/test_helpers.py | 48 ++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 8 deletions(-)
 create mode 100644 gcloud/datastore/test_helpers.py

diff --git a/gcloud/datastore/helpers.py b/gcloud/datastore/helpers.py
index a44424132ea7..ccef721e97b4 100644
--- a/gcloud/datastore/helpers.py
+++ b/gcloud/datastore/helpers.py
@@ -1,5 +1,6 @@
 """Helper methods for dealing with Cloud Datastore's Protobuf API."""
 from datetime import datetime
+import time
 
 import pytz
 
diff --git a/gcloud/datastore/key.py b/gcloud/datastore/key.py
index 13845e1ac270..a02567799b56 100644
--- a/gcloud/datastore/key.py
+++ b/gcloud/datastore/key.py
@@ -49,14 +49,18 @@ def from_protobuf(cls, pb, dataset=None):
   def to_protobuf(self):
     key = datastore_pb.Key()
 
-    # Apparently 's~' is a prefix for High-Replication and is necessary here.
-    # Another valid preflix is 'e~' indicating EU datacenters.
-    dataset_id = self.dataset().id()
-    if dataset_id:
-      if dataset_id[:2] not in ['s~', 'e~']:
-        dataset_id = 's~' + dataset_id
-
-      key.partition_id.dataset_id = dataset_id
+    # Technically a dataset is required to do anything with the key,
+    # but we shouldn't throw a cryptic error if one isn't provided
+    # in the initializer.
+    if self.dataset():
+      # Apparently 's~' is a prefix for High-Replication and is necessary here.
+      # Another valid preflix is 'e~' indicating EU datacenters.
+      dataset_id = self.dataset().id()
+      if dataset_id:
+        if dataset_id[:2] not in ['s~', 'e~']:
+          dataset_id = 's~' + dataset_id
+
+        key.partition_id.dataset_id = dataset_id
 
     if self._namespace:
       key.partition_id.namespace = self._namespace
diff --git a/gcloud/datastore/test_helpers.py b/gcloud/datastore/test_helpers.py
new file mode 100644
index 000000000000..262ebe511a9a
--- /dev/null
+++ b/gcloud/datastore/test_helpers.py
@@ -0,0 +1,48 @@
+from datetime import datetime
+import time
+
+import unittest2
+
+from gcloud.datastore import helpers
+from gcloud.datastore.key import Key
+
+
+class TestHelpers(unittest2.TestCase):
+
+  def test_get_protobuf_attribute(self):
+    mapping = (
+        (str(), 'string_value'),
+        (unicode(), 'string_value'),
+        (int(), 'integer_value'),
+        (long(), 'integer_value'),
+        (float(), 'double_value'),
+        (bool(), 'boolean_value'),
+        (datetime.now(), 'timestamp_microseconds_value'),
+        (Key(), 'key_value'),
+        )
+
+    for test_value, expected_name in mapping:
+      actual_name, _ = helpers.get_protobuf_attribute_and_value(test_value)
+      self.assertEqual(expected_name, actual_name,
+          'Test value "%s" expected %s, got %s' % (
+            test_value, expected_name, actual_name))
+
+  def test_get_protobuf_value(self):
+    now = datetime.now()
+
+    mapping = (
+        (str('string'), 'string'),
+        (unicode('string'), 'string'),
+        (int(), int()),
+        (long(), int()),
+        (float(), float()),
+        (bool(), bool()),
+        (now, time.mktime(now.timetuple())),
+        (Key(), Key().to_protobuf()),
+        )
+
+    for test_value, expected_value in mapping:
+      _, actual_value = helpers.get_protobuf_attribute_and_value(test_value)
+      self.assertEqual(expected_value, actual_value,
+          'Test value "%s" expected %s, got %s.' % (
+            test_value, expected_value, actual_value))