Skip to content

Commit 220170f

Browse files
committed
test coverage: ensure that multi insert preserves job args
1 parent f27a3bb commit 220170f

File tree

1 file changed

+66
-0
lines changed

1 file changed

+66
-0
lines changed

tests/driver/riversqlalchemy/sqlalchemy_driver_test.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,40 @@ async def test_insert_many_tx(self, client, simple_args, test_tx):
306306
assert results[0].job.id > 0
307307

308308

309+
@pytest.mark.asyncio
310+
async def test_insert_many_preserves_distinct_args(self, client):
311+
# Insert mixed types and ensure each row retains its own args and kind
312+
from dataclasses import dataclass
313+
314+
@dataclass
315+
class TypeA:
316+
n: int
317+
kind: str = "simple_a"
318+
319+
def to_json(self) -> str:
320+
return json.dumps({"a": self.n})
321+
322+
@dataclass
323+
class TypeB:
324+
s: str
325+
kind: str = "simple_b"
326+
327+
def to_json(self) -> str:
328+
return json.dumps({"b": self.s})
329+
330+
batch = [TypeA(1), TypeB("x"), TypeA(2), TypeB("y")]
331+
results = await client.insert_many(batch)
332+
333+
assert len(results) == 4
334+
for res, arg in zip(results, batch):
335+
if isinstance(arg, TypeA):
336+
assert res.job.kind == "simple_a"
337+
assert res.job.args == {"a": arg.n}
338+
else:
339+
assert res.job.kind == "simple_b"
340+
assert res.job.args == {"b": arg.s}
341+
342+
309343
class TestSyncClient:
310344
#
311345
# fixtures
@@ -502,3 +536,35 @@ def test_insert_many_tx(self, client, simple_args, test_tx):
502536
assert len(results) == 1
503537
assert results[0].unique_skipped_as_duplicated is False
504538
assert results[0].job.id > 0
539+
540+
def test_insert_many_preserves_distinct_args(self, client):
541+
# Insert mixed types and ensure each row retains its own args and kind
542+
from dataclasses import dataclass
543+
544+
@dataclass
545+
class TypeA:
546+
n: int
547+
kind: str = "simple_a"
548+
549+
def to_json(self) -> str:
550+
return json.dumps({"a": self.n})
551+
552+
@dataclass
553+
class TypeB:
554+
s: str
555+
kind: str = "simple_b"
556+
557+
def to_json(self) -> str:
558+
return json.dumps({"b": self.s})
559+
560+
batch = [TypeA(1), TypeB("x"), TypeA(2), TypeB("y")]
561+
results = client.insert_many(batch)
562+
563+
assert len(results) == 4
564+
for res, arg in zip(results, batch):
565+
if isinstance(arg, TypeA):
566+
assert res.job.kind == "simple_a"
567+
assert res.job.args == {"a": arg.n}
568+
else:
569+
assert res.job.kind == "simple_b"
570+
assert res.job.args == {"b": arg.s}

0 commit comments

Comments
 (0)