Skip to content

Commit 2c9782b

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 ca5c4d8 commit 2c9782b

File tree

2 files changed

+264
-13
lines changed

2 files changed

+264
-13
lines changed

‎tests/test_client_stub.py‎

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,23 +72,121 @@ def test_from_cloud(self):
7272

7373
class TestDgraphClientStubContextManager(helper.ClientIntegrationTestCase):
7474
def setUp(self):
75-
pass
75+
super(TestDgraphClientStubContextManager, self).setUp()
7676

7777
def test_context_manager(self):
78+
"""Test basic context manager usage for DgraphClientStub."""
7879
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
7980
ver = client_stub.check_version(pydgraph.Check())
8081
self.assertIsNotNone(ver)
8182

8283
def test_context_manager_code_exception(self):
84+
"""Test that exceptions within context manager are properly handled."""
8385
with self.assertRaises(AttributeError):
8486
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
8587
self.check_version(client_stub) # AttributeError: no such method
8688

8789
def test_context_manager_function_wrapper(self):
90+
"""Test the client_stub() function wrapper for context manager."""
8891
with pydgraph.client_stub(addr=self.TEST_SERVER_ADDR) as client_stub:
8992
ver = client_stub.check_version(pydgraph.Check())
9093
self.assertIsNotNone(ver)
9194

95+
def test_context_manager_closes_stub(self):
96+
"""Test that the stub is properly closed after exiting context manager."""
97+
stub = None
98+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
99+
stub = client_stub
100+
ver = client_stub.check_version(pydgraph.Check())
101+
self.assertIsNotNone(ver)
102+
103+
# After exiting context, stub should be closed and unusable
104+
with self.assertRaises(Exception):
105+
stub.check_version(pydgraph.Check())
106+
107+
def test_context_manager_with_client(self):
108+
"""Test using DgraphClientStub context manager with DgraphClient."""
109+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
110+
client = pydgraph.DgraphClient(client_stub)
111+
112+
# Perform a simple operation
113+
txn = client.txn(read_only=True)
114+
query = "{ me(func: has(name)) { name } }"
115+
resp = txn.query(query)
116+
self.assertIsNotNone(resp)
117+
118+
def test_context_manager_exception_still_closes(self):
119+
"""Test that stub is closed even when an exception occurs."""
120+
stub_ref = None
121+
try:
122+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
123+
stub_ref = client_stub
124+
client_stub.check_version(pydgraph.Check())
125+
raise ValueError("Test exception")
126+
except ValueError:
127+
pass
128+
129+
# Stub should still be closed despite the exception
130+
with self.assertRaises(Exception):
131+
stub_ref.check_version(pydgraph.Check())
132+
133+
def test_context_manager_function_wrapper_closes(self):
134+
"""Test that client_stub() function wrapper properly closes the stub."""
135+
stub_ref = None
136+
with pydgraph.client_stub(addr=self.TEST_SERVER_ADDR) as client_stub:
137+
stub_ref = client_stub
138+
ver = client_stub.check_version(pydgraph.Check())
139+
self.assertIsNotNone(ver)
140+
141+
# Stub should be closed after exiting
142+
with self.assertRaises(Exception):
143+
stub_ref.check_version(pydgraph.Check())
144+
145+
def test_context_manager_multiple_operations(self):
146+
"""Test performing multiple operations within context manager."""
147+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub:
148+
# Check version multiple times
149+
ver1 = client_stub.check_version(pydgraph.Check())
150+
ver2 = client_stub.check_version(pydgraph.Check())
151+
self.assertIsNotNone(ver1)
152+
self.assertIsNotNone(ver2)
153+
154+
# Create client and perform operations
155+
client = pydgraph.DgraphClient(client_stub)
156+
txn = client.txn(read_only=True)
157+
query = "{ me(func: has(name)) { name } }"
158+
resp = txn.query(query)
159+
self.assertIsNotNone(resp)
160+
161+
def test_context_manager_nested_with_client_operations(self):
162+
"""Test full workflow: stub context manager with client and transaction operations."""
163+
with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as stub:
164+
client = pydgraph.DgraphClient(stub)
165+
166+
# Set schema
167+
schema = "test_name: string @index(fulltext) ."
168+
op = pydgraph.Operation(schema=schema)
169+
client.alter(op)
170+
171+
# Perform mutation and query
172+
with client.txn() as txn:
173+
response = txn.mutate(set_obj={"test_name": "ContextManagerTest"})
174+
self.assertEqual(1, len(response.uids))
175+
uid = list(response.uids.values())[0]
176+
177+
# Verify data was committed
178+
query = '''{{
179+
me(func: uid("{uid}")) {{
180+
test_name
181+
}}
182+
}}'''.format(uid=uid)
183+
184+
with client.txn(read_only=True) as txn:
185+
resp = txn.query(query)
186+
import json
187+
results = json.loads(resp.json).get("me")
188+
self.assertEqual([{"test_name": "ContextManagerTest"}], results)
189+
92190

93191

94192
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)