Added unit tests for tag columns. All tests working.

mdbartos · mdbartos · commit 7e092ad2be8e · 2016-08-29T15:18:43.000-04:00
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -232,6 +232,11 @@ def _convert_dataframe_to_lines(self,
         field_columns = list(field_columns) if list(field_columns) else []
         tag_columns = list(tag_columns) if list(tag_columns) else []
 
+        # If field columns but no tag columns, assume rest of columns are tags
+        if field_columns and (not tag_columns):
+            tag_columns = list(column_series[~column_series.isin(
+                field_columns)])
+
         # Assume that all columns not listed as tag columns are field columns
         if not field_columns:
             field_columns = list(column_series[~column_series.isin(
diff --git a/influxdb/tests/dataframe_client_test.py b/influxdb/tests/dataframe_client_test.py
@@ -69,6 +69,135 @@ def test_write_points_from_dataframe_in_batches(self):
             cli = DataFrameClient(database='db')
             self.assertTrue(cli.write_points(dataframe, "foo", batch_size=1))
 
+    def test_write_points_with_tag_columns(self):
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
+        dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0],
+                                       ['red', 0, "2", 2, 2.0]],
+                                 index=[now, now + timedelta(hours=1)],
+                                 columns=["tag_one", "tag_two", "column_one",
+                                          "column_two", "column_three"])
+        expected = (
+            b"foo,tag_one=blue,tag_two=1 "
+            b"column_one=\"1\",column_two=1i,column_three=1.0 "
+            b"0\n"
+            b"foo,tag_one=red,tag_two=0 "
+            b"column_one=\"2\",column_two=2i,column_three=2.0 "
+            b"3600000000000\n"
+        )
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/write",
+                           status_code=204)
+
+            cli = DataFrameClient(database='db')
+
+            cli.write_points(dataframe, 'foo',
+                             tag_columns=['tag_one', 'tag_two'])
+            self.assertEqual(m.last_request.body, expected)
+
+            cli.write_points(dataframe, 'foo',
+                             tag_columns=['tag_one', 'tag_two'], tags=None)
+            self.assertEqual(m.last_request.body, expected)
+
+    def test_write_points_with_tag_columns_and_global_tags(self):
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
+        dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0],
+                                       ['red', 0, "2", 2, 2.0]],
+                                 index=[now, now + timedelta(hours=1)],
+                                 columns=["tag_one", "tag_two", "column_one",
+                                          "column_two", "column_three"])
+        expected = (
+            b"foo,tag_one=blue,tag_two=1,global_tag=value "
+            b"column_one=\"1\",column_two=1i,column_three=1.0 "
+            b"0\n"
+            b"foo,tag_one=red,tag_two=0,global_tag=value "
+            b"column_one=\"2\",column_two=2i,column_three=2.0 "
+            b"3600000000000\n"
+        )
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/write",
+                           status_code=204)
+
+            cli = DataFrameClient(database='db')
+
+            cli.write_points(dataframe, 'foo',
+                             tag_columns=['tag_one', 'tag_two'],
+                             tags={'global_tag': 'value'})
+            self.assertEqual(m.last_request.body, expected)
+
+    def test_write_points_with_tag_columns_and_defaults(self):
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
+        dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0, 'hot'],
+                                       ['red', 0, "2", 2, 2.0, 'cold']],
+                                 index=[now, now + timedelta(hours=1)],
+                                 columns=["tag_one", "tag_two", "column_one",
+                                          "column_two", "column_three",
+                                          "tag_three"])
+        expected_tags_and_fields = (
+            b"foo,tag_one=blue "
+            b"column_one=\"1\",column_two=1i "
+            b"0\n"
+            b"foo,tag_one=red "
+            b"column_one=\"2\",column_two=2i "
+            b"3600000000000\n"
+        )
+
+        expected_tags_no_fields = (
+            b"foo,tag_one=blue,tag_two=1 "
+            b"column_one=\"1\",column_two=1i,column_three=1.0,"
+            b"tag_three=\"hot\" 0\n"
+            b"foo,tag_one=red,tag_two=0 "
+            b"column_one=\"2\",column_two=2i,column_three=2.0,"
+            b"tag_three=\"cold\" 3600000000000\n"
+        )
+
+        expected_fields_no_tags = (
+            b"foo,tag_one=blue,tag_two=1,tag_three=hot "
+            b"column_one=\"1\",column_two=1i,column_three=1.0 "
+            b"0\n"
+            b"foo,tag_one=red,tag_two=0,tag_three=cold "
+            b"column_one=\"2\",column_two=2i,column_three=2.0 "
+            b"3600000000000\n"
+        )
+
+        expected_no_tags_no_fields = (
+            b"foo "
+            b"tag_one=\"blue\",tag_two=1i,column_one=\"1\","
+            b"column_two=1i,column_three=1.0,tag_three=\"hot\" "
+            b"0\n"
+            b"foo "
+            b"tag_one=\"red\",tag_two=0i,column_one=\"2\","
+            b"column_two=2i,column_three=2.0,tag_three=\"cold\" "
+            b"3600000000000\n"
+        )
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/write",
+                           status_code=204)
+
+            cli = DataFrameClient(database='db')
+
+            cli.write_points(dataframe, 'foo',
+                             field_columns=['column_one', 'column_two'],
+                             tag_columns=['tag_one'])
+            self.assertEqual(m.last_request.body, expected_tags_and_fields)
+
+            cli.write_points(dataframe, 'foo',
+                             tag_columns=['tag_one', 'tag_two'])
+            self.assertEqual(m.last_request.body, expected_tags_no_fields)
+
+            cli.write_points(dataframe, 'foo',
+                             field_columns=['column_one', 'column_two',
+                                            'column_three'])
+            self.assertEqual(m.last_request.body, expected_fields_no_tags)
+
+            cli.write_points(dataframe, 'foo')
+            self.assertEqual(m.last_request.body, expected_no_tags_no_fields)
+
     def test_write_points_from_dataframe_with_numeric_column_names(self):
         now = pd.Timestamp('1970-01-01 00:00+00:00')
         # df with numeric column names
