fix: creating Tasks with import statements (influxdata#491)

bednar · web-flow · commit fade9e08ca6e · 2022-08-15T13:24:52.000-06:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,7 +5,8 @@
 ### Bug Fixes
 1. [#483](https://github.com/influxdata/influxdb-client-python/pull/483): Querying data if the `debug` is enabled
 1. [#477](https://github.com/influxdata/influxdb-client-python/pull/477): Parsing date fails due to thread race
-1. [#486](https://github.com/influxdata/influxdb-client-python/pull/486): Fix bug when serializing DataFrames that might occur if you're inserting NaN values and have columns starting with digits.
+1. [#486](https://github.com/influxdata/influxdb-client-python/pull/486): Serializing DataFrames with columns starting with digits
+1. [#491](https://github.com/influxdata/influxdb-client-python/pull/491): Creating `Tasks` with `import` statements
 
 ### Dependencies
 1. [#472](https://github.com/influxdata/influxdb-client-python/pull/472): Update `RxPY` to `4.0.4`
diff --git a/influxdb_client/client/tasks_api.py b/influxdb_client/client/tasks_api.py
@@ -4,7 +4,6 @@
 Use tasks (scheduled Flux queries) to input a data stream and then analyze, modify, and act on the data accordingly.
 """
 
-
 import datetime
 from typing import List
 
@@ -65,7 +64,7 @@ def _create_task(name: str, flux: str, every, cron, org_id: str) -> Task:
             repetition += "cron: "
             repetition += '"' + cron + '"'
 
-        flux_with_options = 'option task = {{name: "{}", {}}} \n {}'.format(name, repetition, flux)
+        flux_with_options = '{} \n\noption task = {{name: "{}", {}}}'.format(flux, name, repetition)
         task.flux = flux_with_options
 
         return task
@@ -151,10 +150,10 @@ def get_runs(self, task_id, **kwargs) -> List['Run']:
         Retrieve list of run records for a task.
 
         :param task_id: task id
-        :param str after: returns runs after specified ID
-        :param int limit: the number of runs to return
-        :param datetime after_time: filter runs to those scheduled after this time, RFC3339
-        :param datetime before_time: filter runs to those scheduled before this time, RFC3339
+        :key str after: returns runs after specified ID
+        :key int limit: the number of runs to return
+        :key datetime after_time: filter runs to those scheduled after this time, RFC3339
+        :key datetime before_time: filter runs to those scheduled before this time, RFC3339
         """
         return self._service.get_tasks_id_runs(task_id=task_id, **kwargs).runs
 
diff --git a/tests/test_TasksApi.py b/tests/test_TasksApi.py
@@ -121,7 +121,7 @@ def test_create_task_every(self):
         self.assertEqual(task.status, "active")
         self.assertEqual(task.every, "1h")
         self.assertEqual(task.cron, None)
-        self.assertTrue(task.flux.endswith(TASK_FLUX))
+        self.assertTrue(task.flux.startswith(TASK_FLUX))
 
     def test_create_task_cron(self):
         task_name = self.generate_name("it task")
@@ -137,7 +137,7 @@ def test_create_task_cron(self):
         self.assertEqual(task.cron, "0 2 * * *")
         # self.assertEqualIgnoringWhitespace(task.flux, flux)
 
-        self.assertTrue(task.flux.endswith(TASK_FLUX))
+        self.assertTrue(task.flux.startswith(TASK_FLUX))
         # self.assertEqual(task.links, "active")
 
         links = task.links
@@ -151,6 +151,23 @@ def test_create_task_cron(self):
         # TODO missing get labels
         self.assertEqual(links.labels, "/api/v2/tasks/" + task.id + "/labels")
 
+    def test_create_with_import(self):
+        task_name = self.generate_name("it task")
+        task_flux = 'import "http"\n\n' \
+                    'from(bucket: "iot_center")\n' \
+                    '    |> range(start: -30d)\n' \
+                    '    |> filter(fn: (r) => r._measurement == "environment")\n' \
+                    '    |> aggregateWindow(every: 1h, fn: mean)'
+        task = self.tasks_api.create_task_cron(task_name, task_flux, "10 0 * * * *", self.organization.id)
+
+        self.assertIsNotNone(task.id)
+        self.assertEqual(task.name, task_name)
+        self.assertEqual(task.org_id, self.organization.id)
+        self.assertEqual(task.status, "active")
+        self.assertEqual(task.cron, "10 0 * * * *")
+        self.assertTrue(task.flux.startswith(task_flux))
+        self.assertTrue(task.flux.splitlines()[-1].startswith('option task = '))
+
     def test_find_task_by_id(self):
         task_name = self.generate_name("it task")
         task = self.tasks_api.create_task_cron(task_name, TASK_FLUX, "0 2 * * *", self.organization.id)
@@ -182,12 +199,12 @@ def test_update_task(self):
         cron_task = self.tasks_api.create_task_cron(task_name, TASK_FLUX, "0 2 * * *", self.organization.id)
 
         flux = '''
+        {flux}
+        
         option task = {{
             name: "{task_name}",
             every: 3m
         }}
-        
-        {flux}
         '''.format(task_name=task_name, flux=TASK_FLUX)
 
         cron_task.cron = None
