tests/artifactcache: Use context manager for gRPC channels

1 files changed, 50 insertions, 50 deletions

diff --git a/tests/artifactcache/artifactservice.py b/tests/artifactcache/artifactservice.py
index b3bc7c218..00d14b45d 100644
--- a/tests/artifactcache/artifactservice.py
+++ b/tests/artifactcache/artifactservice.py
@@ -38,19 +38,19 @@ def test_artifact_get_not_found(tmpdir):
     with create_artifact_share(sharedir) as share:
         # set up artifact service stub
         url = urlparse(share.repo)
-        channel = grpc.insecure_channel("{}:{}".format(url.hostname, url.port))
-        artifact_stub = ArtifactServiceStub(channel)
+        with grpc.insecure_channel("{}:{}".format(url.hostname, url.port)) as channel:
+            artifact_stub = ArtifactServiceStub(channel)
 
-        # Run GetArtifact and check it throws a not found error
-        request = GetArtifactRequest()
-        request.cache_key = "@artifact/something/not_there"
-        try:
-            artifact_stub.GetArtifact(request)
-        except grpc.RpcError as e:
-            assert e.code() == grpc.StatusCode.NOT_FOUND
-            assert e.details() == "Artifact proto not found"
-        else:
-            assert False
+            # Run GetArtifact and check it throws a not found error
+            request = GetArtifactRequest()
+            request.cache_key = "@artifact/something/not_there"
+            try:
+                artifact_stub.GetArtifact(request)
+            except grpc.RpcError as e:
+                assert e.code() == grpc.StatusCode.NOT_FOUND
+                assert e.details() == "Artifact proto not found"
+            else:
+                assert False
 
 
 # Successfully getting the artifact
@@ -70,42 +70,42 @@ def test_update_artifact(tmpdir, files):
 
         url = urlparse(share.repo)
 
-        channel = grpc.insecure_channel("{}:{}".format(url.hostname, url.port))
-        artifact_stub = ArtifactServiceStub(channel)
-
-        # initialise an artifact
-        artifact = Artifact()
-        artifact.version = 0
-        artifact.build_success = True
-        artifact.strong_key = "abcdefghijklmnop"
-        artifact.files.hash = "hashymchashash"
-        artifact.files.size_bytes = 10
-
-        artifact.files.CopyFrom(digest)
-
-        # Put it in the artifact share with an UpdateArtifactRequest
-        request = UpdateArtifactRequest()
-        request.artifact.CopyFrom(artifact)
-        request.cache_key = "a-cache-key"
-
-        # should return the same artifact back
-        if files == "present":
-            response = artifact_stub.UpdateArtifact(request)
+        with grpc.insecure_channel("{}:{}".format(url.hostname, url.port)) as channel:
+            artifact_stub = ArtifactServiceStub(channel)
+
+            # initialise an artifact
+            artifact = Artifact()
+            artifact.version = 0
+            artifact.build_success = True
+            artifact.strong_key = "abcdefghijklmnop"
+            artifact.files.hash = "hashymchashash"
+            artifact.files.size_bytes = 10
+
+            artifact.files.CopyFrom(digest)
+
+            # Put it in the artifact share with an UpdateArtifactRequest
+            request = UpdateArtifactRequest()
+            request.artifact.CopyFrom(artifact)
+            request.cache_key = "a-cache-key"
+
+            # should return the same artifact back
+            if files == "present":
+                response = artifact_stub.UpdateArtifact(request)
+                assert response == artifact
+            else:
+                try:
+                    artifact_stub.UpdateArtifact(request)
+                except grpc.RpcError as e:
+                    assert e.code() == grpc.StatusCode.FAILED_PRECONDITION
+                    if files == "absent":
+                        assert e.details() == "Artifact files specified but no files found"
+                    elif files == "invalid":
+                        assert e.details() == "Artifact files specified but directory not found"
+                    return
+
+            # If we uploaded the artifact check GetArtifact
+            request = GetArtifactRequest()
+            request.cache_key = "a-cache-key"
+
+            response = artifact_stub.GetArtifact(request)
             assert response == artifact
-        else:
-            try:
-                artifact_stub.UpdateArtifact(request)
-            except grpc.RpcError as e:
-                assert e.code() == grpc.StatusCode.FAILED_PRECONDITION
-                if files == "absent":
-                    assert e.details() == "Artifact files specified but no files found"
-                elif files == "invalid":
-                    assert e.details() == "Artifact files specified but directory not found"
-                return
-
-        # If we uploaded the artifact check GetArtifact
-        request = GetArtifactRequest()
-        request.cache_key = "a-cache-key"
-
-        response = artifact_stub.GetArtifact(request)
-        assert response == artifact