diff options
-rw-r--r-- | Lib/test/test_urllib_response.py | 42 | ||||
-rw-r--r-- | Lib/urllib/response.py | 8 |
2 files changed, 50 insertions, 0 deletions
diff --git a/Lib/test/test_urllib_response.py b/Lib/test/test_urllib_response.py new file mode 100644 index 0000000..fdd3325 --- /dev/null +++ b/Lib/test/test_urllib_response.py @@ -0,0 +1,42 @@ +"""Unit tests for code in urllib.response.""" + +import test.support +import urllib.response +import unittest + +class TestFile(object): + + def __init__(self): + self.closed = False + + def read(self, bytes): + pass + + def readline(self): + pass + + def close(self): + self.closed = True + +class Testaddbase(unittest.TestCase): + + # TODO(jhylton): Write tests for other functionality of addbase() + + def setUp(self): + self.fp = TestFile() + self.addbase = urllib.response.addbase(self.fp) + + def test_with(self): + def f(): + with self.addbase as spam: + pass + self.assertFalse(self.fp.closed) + f() + self.assertTrue(self.fp.closed) + self.assertRaises(ValueError, f) + +def test_main(): + test.support.run_unittest(Testaddbase) + +if __name__ == '__main__': + test_main() diff --git a/Lib/urllib/response.py b/Lib/urllib/response.py index 52eeed0..9859642 100644 --- a/Lib/urllib/response.py +++ b/Lib/urllib/response.py @@ -40,6 +40,14 @@ class addbase(object): if self.fp: self.fp.close() self.fp = None + def __enter__(self): + if self.fp is None: + raise ValueError("I/O operation on closed file") + return self + + def __exit__(self, type, value, traceback): + self.close() + class addclosehook(addbase): """Class to add a close hook to an open file.""" |