djbdns/io.py: fix test that fails in other timezones.

[djbdns-logparse.git] / djbdns / io.py

diff --git a/djbdns/io.py b/djbdns/io.py
index 7c70bab3debcf2b78ad43e18146488740cdf1c4e..c3952bed7da220119a68604e2ed6d0f26d098114 100644 (file)
--- a/djbdns/io.py
+++ b/djbdns/io.py
@@ -1,3 +1,7 @@
+r"""
+Functions that perform input/output. This forms a layer between the
+executable itself and the more libraryish modules.
+"""
 from subprocess import Popen, PIPE
 from typing import TextIO
 from djbdns.dnscache import handle_dnscache_log
@@ -20,6 +24,8 @@ def parse_logfile(file : TextIO):
         >>> from tempfile import NamedTemporaryFile
         >>> with NamedTemporaryFile(mode="w", delete=False) as f:
         ...     _ = f.write(line)
+        >>> from os import environ
+        >>> environ["TZ"] = "UTC+4"
         >>> f = open(f.name, 'r')
         >>> parse_logfile(f)
         2022-09-14 21:04:40.206516500 dropped query (no authority) from 127.0.0.1:40289 (id 48745): a www.example.com
@@ -31,19 +37,31 @@ def parse_logfile(file : TextIO):
     # Open a pipe to tai64nlocal. We'll write lines of our input file
     # (the log file) to it, and read back the same lines but with
     # friendly timestamps in them.
-    tai = Popen(["tai64nlocal"], stdin=PIPE, stdout=PIPE, text=True, bufsize=0)
+    with Popen(["tai64nlocal"],
+               stdin=PIPE,
+               stdout=PIPE,
+               text=True,
+               bufsize=0) as tai:
 
-    for line in file:
-        tai.stdin.write(line)
-        line = tai.stdout.readline()
+        if not tai.stdin or not tai.stdout:
+            # Mypy tells me that this can happen, based on the type
+            # annotations in the standard library I guess?
+            return
 
-        friendly_line = handle_tinydns_log(line)
-        if not friendly_line:
-            friendly_line = handle_dnscache_log(line)
-            if not friendly_line:
-                friendly_line = line
+        for line in file:
+            tai.stdin.write(line)
+            line = tai.stdout.readline()
 
-        print(friendly_line)
+            friendly_line = handle_tinydns_log(line)
+            if not friendly_line:
+                friendly_line = handle_dnscache_log(line)
+                if not friendly_line:
+                    friendly_line = line
 
-    # Ensure that the pipe gets closed.
-    tai.communicate()
+            try:
+                print(friendly_line)
+            except BrokenPipeError:
+                # If our stdout is being piped to another process and if
+                # that process closes the pipe, this error will be raised
+                # the next time we try to write to stdout.
+                break