Skip to content

Commit 2fc50e4

Browse files
committed
Add a couple more enum tests
1 parent ffd134e commit 2fc50e4

File tree

1 file changed

+39
-0
lines changed

1 file changed

+39
-0
lines changed

tests/test_codecs.py

+39
Original file line numberDiff line numberDiff line change
@@ -1624,6 +1624,45 @@ async def test_enum_and_range(self):
16241624
DROP TYPE enum_t;
16251625
''')
16261626

1627+
async def test_enum_in_composite(self):
1628+
await self.con.execute('''
1629+
CREATE TYPE enum_t AS ENUM ('abc', 'def', 'ghi');
1630+
CREATE TYPE composite_w_enum AS (a int, b enum_t);
1631+
''')
1632+
1633+
try:
1634+
result = await self.con.fetchval('''
1635+
SELECT ROW(1, 'def'::enum_t)::composite_w_enum
1636+
''')
1637+
self.assertEqual(set(result.items()), {('a', 1), ('b', 'def')})
1638+
1639+
finally:
1640+
await self.con.execute('''
1641+
DROP TYPE composite_w_enum;
1642+
DROP TYPE enum_t;
1643+
''')
1644+
1645+
async def test_enum_function_return(self):
1646+
await self.con.execute('''
1647+
CREATE TYPE enum_t AS ENUM ('abc', 'def', 'ghi');
1648+
CREATE FUNCTION return_enum() RETURNS enum_t
1649+
LANGUAGE plpgsql AS $$
1650+
BEGIN
1651+
RETURN 'abc'::enum_t;
1652+
END;
1653+
$$;
1654+
''')
1655+
1656+
try:
1657+
result = await self.con.fetchval('''SELECT return_enum()''')
1658+
self.assertEqual(result, 'abc')
1659+
1660+
finally:
1661+
await self.con.execute('''
1662+
DROP FUNCTION return_enum();
1663+
DROP TYPE enum_t;
1664+
''')
1665+
16271666
async def test_no_result(self):
16281667
st = await self.con.prepare('rollback')
16291668
self.assertTupleEqual(st.get_attributes(), ())

0 commit comments

Comments
 (0)