Minor refactoring the asyncnotify test.

James William Pye · James William Pye · commit e386cdaa590e · 2010-03-27T20:14:02.000-07:00
The problem was that the CopyFail was not occuring in some cases. So make the
test avoid potential failure(likely due to OS buffer size differences).

Additionally, add a test that shows early breaks being detected by the
CopyManager.
diff --git a/postgresql/copyman.py b/postgresql/copyman.py
@@ -676,6 +676,7 @@ def __init__(self, producer, *receivers):
 		self.producer = producer
 		self.transformer = None
 		self.receivers = ElementSet(receivers)
+		self._seen_stop_iteration = False
 		rp = set()
 		add = rp.add
 		for x in self.receivers:
@@ -739,7 +740,7 @@ def __exit__(self, typ, val, tb):
 			except Exception as e:
 				exit_faults[x] = e
 
-		if typ or exit_faults or profail:
+		if typ or exit_faults or profail or not self._seen_stop_iteration:
 			raise CopyFail(self,
 				"could not complete the COPY operation",
 				receiver_faults = exit_faults or None,
@@ -770,6 +771,7 @@ def _service_producer(self):
 			nextdata = next(self.producer)
 		except StopIteration:
 			# Should be over.
+			self._seen_stop_iteration = True
 			raise
 		except Exception:
 			raise ProducerFault(self)
diff --git a/postgresql/test/test_copyman.py b/postgresql/test/test_copyman.py
@@ -247,19 +247,32 @@ def testAsyncNotify(self):
 		sqlexec(stdsource)
 		dst = new()
 		dst.execute(stddst)
-		stmt = prepare(srcsql)
-		sp = Injector(notify, prepare(srcsql), buffer_size = 133)
+		sp = Injector(notify, prepare(srcsql), buffer_size = 32)
 		sr = copyman.StatementReceiver(dst.prepare(dstsql))
 		seen_in_loop = 0
+		r = []
 		with copyman.CopyManager(sp, sr) as copy:
 			for x in copy:
-				r = list(db.iternotifies(0))
-				if r:
-					break
-			else:
-				self.fail("didn't pickup notify during copy")
+				r += list(db.iternotifies(0))
 		# Got the injected NOTIFY's, right?
-		self.failUnlessEqual(r, [('channel', 'payload', 1234)])
+		self.failUnless(r)
+		# it may have happened multiple times, so adjust accordingly.
+		self.failUnlessEqual(r, [('channel', 'payload', 1234)]*len(r))
+
+	@pg_tmp
+	def testUnfinishedCopy(self):
+		sqlexec(stdsource)
+		dst = new()
+		dst.execute(stddst)
+		sp = copyman.StatementProducer(prepare(srcsql), buffer_size = 32)
+		sr = copyman.StatementReceiver(dst.prepare(dstsql))
+		try:
+			with copyman.CopyManager(sp, sr) as copy:
+				for x in copy:
+					break
+			self.fail("did not raise CopyFail")
+		except copyman.CopyFail:
+			pass
 
 	@pg_tmp
 	def testRaiseInCopy(self):
