Skip to content

Commit cb32f56

Browse files
shaunpattersonclaude
shaunpatterson
and
claude
committed
Add comprehensive tests for context manager functionality
This commit adds extensive test coverage for the context manager implementation in PR #185, which adds Python context manager support to the Txn and DgraphClientStub classes. Tests for Txn context manager: - Auto-commit for write transactions on successful completion - Auto-discard for read-only transactions - Exception handling with automatic discard - Transaction state validation after context manager exit - Multiple mutations within a single context manager - Query and mutate operations combined - Invalid operation exception handling - Read-only transaction mutation prevention - No-mutation transaction handling Tests for DgraphClientStub context manager: - Basic context manager usage - Exception handling - Proper resource cleanup (stub close) after exit - Exception-safe resource cleanup - Multiple operations within context manager - Integration with DgraphClient and transactions - Full workflow with nested context managers These tests ensure robust behavior of the context manager implementation including proper resource cleanup, exception handling, and transaction lifecycle management. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 <[email protected]>
1 parent fe8d9b5 commit cb32f56

File tree

2 files changed

+267
-16
lines changed

2 files changed

+267
-16
lines changed

‎tests/test_client_stub.py‎

Lines changed: 102 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,22 +70,120 @@ def test_from_cloud(self):
7070

7171
class TestDgraphClientStubContextManager(helper.ClientIntegrationTestCase):
7272
def setUp(self):
73-
pass
74-
73+
super(TestDgraphClientStubContextManager, self).setUp()
74+
7575
def test_context_manager(self):
76+
"""Test basic context manager usage for DgraphClientStub."""
7677
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
7778
ver = client_stub.check_version(pydgraph.Check())
7879
self.assertIsNotNone(ver)
79-
80+
8081
def test_context_manager_code_exception(self):
82+
"""Test that exceptions within context manager are properly handled."""
8183
with self.assertRaises(AttributeError):
8284
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
8385
self.check_version(client_stub) # AttributeError: no such method
84-
86+
8587
def test_context_manager_function_wrapper(self):
88+
"""Test the client_stub() function wrapper for context manager."""
89+
with pydgraph.client_stub(addr=self.TEST_SERVER_ADDR) as client_stub:
90+
ver = client_stub.check_version(pydgraph.Check())
91+
self.assertIsNotNone(ver)
92+
93+
def test_context_manager_closes_stub(self):
94+
"""Test that the stub is properly closed after exiting context manager."""
95+
stub = None
96+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
97+
stub = client_stub
98+
ver = client_stub.check_version(pydgraph.Check())
99+
self.assertIsNotNone(ver)
100+
101+
# After exiting context, stub should be closed and unusable
102+
with self.assertRaises(Exception):
103+
stub.check_version(pydgraph.Check())
104+
105+
def test_context_manager_with_client(self):
106+
"""Test using DgraphClientStub context manager with DgraphClient."""
107+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
108+
client = pydgraph.DgraphClient(client_stub)
109+
110+
# Perform a simple operation
111+
txn = client.txn(read_only=True)
112+
query = "{ me(func: has(name)) { name } }"
113+
resp = txn.query(query)
114+
self.assertIsNotNone(resp)
115+
116+
def test_context_manager_exception_still_closes(self):
117+
"""Test that stub is closed even when an exception occurs."""
118+
stub_ref = None
119+
try:
120+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
121+
stub_ref = client_stub
122+
client_stub.check_version(pydgraph.Check())
123+
raise ValueError("Test exception")
124+
except ValueError:
125+
pass
126+
127+
# Stub should still be closed despite the exception
128+
with self.assertRaises(Exception):
129+
stub_ref.check_version(pydgraph.Check())
130+
131+
def test_context_manager_function_wrapper_closes(self):
132+
"""Test that client_stub() function wrapper properly closes the stub."""
133+
stub_ref = None
86134
with pydgraph.client_stub(addr=self.TEST_SERVER_ADDR) as client_stub:
135+
stub_ref = client_stub
87136
ver = client_stub.check_version(pydgraph.Check())
88137
self.assertIsNotNone(ver)
138+
139+
# Stub should be closed after exiting
140+
with self.assertRaises(Exception):
141+
stub_ref.check_version(pydgraph.Check())
142+
143+
def test_context_manager_multiple_operations(self):
144+
"""Test performing multiple operations within context manager."""
145+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
146+
# Check version multiple times
147+
ver1 = client_stub.check_version(pydgraph.Check())
148+
ver2 = client_stub.check_version(pydgraph.Check())
149+
self.assertIsNotNone(ver1)
150+
self.assertIsNotNone(ver2)
151+
152+
# Create client and perform operations
153+
client = pydgraph.DgraphClient(client_stub)
154+
txn = client.txn(read_only=True)
155+
query = "{ me(func: has(name)) { name } }"
156+
resp = txn.query(query)
157+
self.assertIsNotNone(resp)
158+
159+
def test_context_manager_nested_with_client_operations(self):
160+
"""Test full workflow: stub context manager with client and transaction operations."""
161+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as stub:
162+
client = pydgraph.DgraphClient(stub)
163+
164+
# Set schema
165+
schema = "test_name: string @index(fulltext) ."
166+
op = pydgraph.Operation(schema=schema)
167+
client.alter(op)
168+
169+
# Perform mutation and query
170+
with client.txn() as txn:
171+
response = txn.mutate(set_obj={"test_name": "ContextManagerTest"})
172+
self.assertEqual(1, len(response.uids))
173+
uid = list(response.uids.values())[0]
174+
175+
# Verify data was committed
176+
query = '''{{
177+
me(func: uid("{uid}")) {{
178+
test_name
179+
}}
180+
}}'''.format(uid=uid)
181+
182+
with client.txn(read_only=True) as txn:
183+
resp = txn.query(query)
184+
import json
185+
results = json.loads(resp.json).get("me")
186+
self.assertEqual([{"test_name": "ContextManagerTest"}], results)
89187

90188

91189
def suite():

‎tests/test_txn.py‎

Lines changed: 165 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -612,31 +612,184 @@ def test_sp_star2(self):
612612

613613
class TestContextManager(helper.ClientIntegrationTestCase):
614614
def setUp(self):
615-
self.stub = pydgraph.DgraphClientStub(self.TEST_SERVER_ADDR)
616-
self.client = pydgraph.DgraphClient(self.stub)
617-
self.q = '''
615+
super(TestContextManager, self).setUp()
616+
helper.drop_all(self.client)
617+
helper.set_schema(self.client, "name: string @index(fulltext) .")
618+
619+
def test_context_manager_by_contextlib(self):
620+
"""Test context manager via client.begin() for read-only transactions."""
621+
q = '''
618622
{
619623
company(func: type(x.Company), first: 10){
620624
expand(_all_)
621625
}
622626
}
623627
'''
624-
def tearDown(self) -> None:
625-
self.stub.close()
626-
627-
def test_context_manager_by_contextlib(self):
628628
with self.client.begin(read_only=True, best_effort=True) as tx:
629-
response = tx.query(self.q)
629+
response = tx.query(q)
630630
self.assertIsNotNone(response)
631631
data = json.loads(response.json)
632-
print(data)
633-
632+
634633
def test_context_manager_by_class(self):
634+
"""Test context manager using Txn class directly for read-only transactions."""
635+
q = '''
636+
{
637+
company(func: type(x.Company), first: 10){
638+
expand(_all_)
639+
}
640+
}
641+
'''
635642
with pydgraph.Txn(self.client, read_only=True, best_effort=True) as tx:
636-
response = tx.query(self.q)
643+
response = tx.query(q)
637644
self.assertIsNotNone(response)
638645
data = json.loads(response.json)
639-
print(data)
646+
647+
def test_context_manager_auto_commit(self):
648+
"""Test that write transactions automatically commit on successful completion."""
649+
with self.client.txn() as txn:
650+
response = txn.mutate(set_obj={"name": "Alice"})
651+
self.assertEqual(1, len(response.uids), "Nothing was assigned")
652+
uid = list(response.uids.values())[0]
653+
654+
# Verify the data was committed by querying in a new transaction
655+
query = '''{{
656+
me(func: uid("{uid}")) {{
657+
name
658+
}}
659+
}}'''.format(uid=uid)
660+
661+
resp = self.client.txn(read_only=True).query(query)
662+
self.assertEqual([{"name": "Alice"}], json.loads(resp.json).get("me"))
663+
664+
def test_context_manager_read_only_auto_discard(self):
665+
"""Test that read-only transactions automatically discard."""
666+
# Create some data first
667+
txn = self.client.txn()
668+
response = txn.mutate(set_obj={"name": "Bob"})
669+
uid = list(response.uids.values())[0]
670+
txn.commit()
671+
672+
# Read-only transaction should auto-discard (not commit)
673+
query = '''{{
674+
me(func: uid("{uid}")) {{
675+
name
676+
}}
677+
}}'''.format(uid=uid)
678+
679+
with self.client.txn(read_only=True) as txn:
680+
resp = txn.query(query)
681+
self.assertEqual([{"name": "Bob"}], json.loads(resp.json).get("me"))
682+
683+
# Transaction should be finished after context manager exits
684+
self.assertTrue(txn._finished)
685+
686+
def test_context_manager_exception_handling(self):
687+
"""Test that exceptions cause automatic discard and are re-raised."""
688+
with self.assertRaises(ValueError):
689+
with self.client.txn() as txn:
690+
response = txn.mutate(set_obj={"name": "Charlie"})
691+
uid = list(response.uids.values())[0]
692+
raise ValueError("Test exception")
693+
694+
# Verify transaction was discarded - data should not exist
695+
query = '''{{
696+
me(func: has(name)) {{
697+
name
698+
}}
699+
}}'''
700+
701+
resp = self.client.txn(read_only=True).query(query)
702+
results = json.loads(resp.json).get("me")
703+
# Should be empty or not contain Charlie
704+
if results:
705+
names = [r.get("name") for r in results]
706+
self.assertNotIn("Charlie", names)
707+
708+
def test_context_manager_transaction_finished_after_exit(self):
709+
"""Test that transaction is marked as finished after exiting context manager."""
710+
with self.client.txn() as txn:
711+
txn.mutate(set_obj={"name": "David"})
712+
self.assertFalse(txn._finished)
713+
714+
# Should be finished after exit
715+
self.assertTrue(txn._finished)
716+
717+
# Should not be able to use transaction after context manager
718+
with self.assertRaises(Exception):
719+
txn.query("{ me() {} }")
720+
721+
def test_context_manager_multiple_mutations(self):
722+
"""Test multiple mutations within a single context manager."""
723+
with self.client.txn() as txn:
724+
response1 = txn.mutate(set_obj={"name": "Eve"})
725+
uid1 = list(response1.uids.values())[0]
726+
727+
response2 = txn.mutate(set_obj={"name": "Frank"})
728+
uid2 = list(response2.uids.values())[0]
729+
730+
# Verify both mutations were committed
731+
query = '''{{
732+
me(func: has(name), orderasc: name) {{
733+
name
734+
}}
735+
}}'''
736+
737+
resp = self.client.txn(read_only=True).query(query)
738+
results = json.loads(resp.json).get("me")
739+
names = [r.get("name") for r in results]
740+
self.assertIn("Eve", names)
741+
self.assertIn("Frank", names)
742+
743+
def test_context_manager_query_and_mutate(self):
744+
"""Test both query and mutate operations within a context manager."""
745+
# Create initial data
746+
txn = self.client.txn()
747+
response = txn.mutate(set_obj={"name": "Grace"})
748+
uid = list(response.uids.values())[0]
749+
txn.commit()
750+
751+
# Query and update in context manager
752+
with self.client.txn() as txn:
753+
query = '''{{
754+
me(func: uid("{uid}")) {{
755+
name
756+
}}
757+
}}'''.format(uid=uid)
758+
759+
resp = txn.query(query)
760+
self.assertEqual([{"name": "Grace"}], json.loads(resp.json).get("me"))
761+
762+
# Update the name
763+
txn.mutate(set_obj={"uid": uid, "name": "Grace Updated"})
764+
765+
# Verify the update was committed
766+
resp = self.client.txn(read_only=True).query(query)
767+
self.assertEqual([{"name": "Grace Updated"}], json.loads(resp.json).get("me"))
768+
769+
def test_context_manager_invalid_nquad_exception(self):
770+
"""Test that invalid operations cause proper exception handling and discard."""
771+
with self.assertRaises(Exception):
772+
with self.client.txn() as txn:
773+
# This should fail with invalid N-Quad syntax
774+
txn.mutate(set_nquads="_:node <name> InvalidWithoutQuotes")
775+
776+
# Transaction should be finished
777+
self.assertTrue(txn._finished)
778+
779+
def test_context_manager_read_only_cannot_mutate(self):
780+
"""Test that read-only transactions cannot mutate within context manager."""
781+
with self.assertRaises(Exception):
782+
with self.client.txn(read_only=True) as txn:
783+
txn.mutate(set_obj={"name": "Should Fail"})
784+
785+
def test_context_manager_no_mutations_auto_commit(self):
786+
"""Test that transactions with no mutations don't error on auto-commit."""
787+
with self.client.txn() as txn:
788+
query = "{ me(func: has(name)) { name } }"
789+
resp = txn.query(query)
790+
791+
# Should complete without errors even though no mutations were made
792+
self.assertTrue(txn._finished)
640793

641794

642795
def suite():

0 commit comments

Comments
 (0)