diff --git a/src/infi/clickhouse_orm/database.py b/src/infi/clickhouse_orm/database.py index 703e982..9435946 100644 --- a/src/infi/clickhouse_orm/database.py +++ b/src/infi/clickhouse_orm/database.py @@ -245,7 +245,7 @@ class Database(object): # Return any remaining lines in partial batch if lines: yield buf.getvalue() - self._send(gen()) + return self._send(gen()) def count(self, model_class, conditions=None): ''' diff --git a/tests/test_database.py b/tests/test_database.py index 44971e1..96b47da 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -46,6 +46,11 @@ class DatabaseTestCase(TestCaseWithData): self.assertEqual(str(t.b), '2020-01-01') self.assertEqual(t.d, 35) + def test_insert__response(self): + from requests import Response + response = self.database.insert(next(self._sample_data())) + self.assertTrue(isinstance(response, Response)) + def test_count(self): self.database.insert(self._sample_data()) self.assertEqual(self.database.count(Person), 100)