fix: flaky test

Author: NguyenHoangSon96

auto_test.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+total_count=0
+error_count=0
+
+for (( i = 0; i < 100; i++ )); do
+  result=$(pytest tests/test_polars.py::TestWritePolars::test_write_polars_batching)
+  if [[ "$result" == *"expected call not found"* ]]; then
+    echo "$result"
+    ((error_count++))
+  fi
+
+  ((total_count++))
+done
+
+echo "Total run $total_count"
+echo "Error count $error_count"

tests/test_polars.py

@@ -1,6 +1,7 @@
 import importlib.util
+import time
 import unittest
-from unittest.mock import Mock
+from unittest.mock import Mock, ANY
 
 from influxdb_client_3 import PointSettings, InfluxDBClient3, write_client_options, WriteOptions
 from influxdb_client_3.write_client import WriteService
@@ -86,6 +87,12 @@ def test_write_polars_batching(self):
             data_frame_timestamp_column="time",
         )
 
-        actual = self.client._write_api._write_service.post_write.call_args[1]['body']
-        self.assertEqual(b'measurement temperature=22.4 1722470400000000000\n'
-                         b'measurement temperature=21.8 1722474000000000000', actual)
+        time.sleep(0.5)
+        self.client._write_api._write_service.post_write.assert_called_once_with(
+            org=ANY,
+            bucket=ANY,
+            precision=ANY,
+            async_req=ANY,
+            content_type=ANY,
+            urlopen_kw=ANY,
+            body=b'measurement temperature=22.4 1722470400000000000\nmeasurement temperature=21.8 1722474000000000000')