Improve e2e test traffic assertion

Author: albin-mullvad

android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/LeakTest.kt

@@ -15,8 +15,9 @@ import net.mullvad.mullvadvpn.test.common.page.on
 import net.mullvad.mullvadvpn.test.common.rule.ForgetAllVpnAppsInSettingsTestRule
 import net.mullvad.mullvadvpn.test.e2e.annotations.HasDependencyOnLocalAPI
 import net.mullvad.mullvadvpn.test.e2e.misc.AccountTestRule
-import net.mullvad.mullvadvpn.test.e2e.misc.LeakCheck
+import net.mullvad.mullvadvpn.test.e2e.misc.NetworkTrafficChecker
 import net.mullvad.mullvadvpn.test.e2e.misc.NoTrafficToHostRule
+import net.mullvad.mullvadvpn.test.e2e.misc.OnlyTrafficToHostRule
 import net.mullvad.mullvadvpn.test.e2e.misc.TrafficGenerator
 import net.mullvad.mullvadvpn.test.e2e.router.packetCapture.PacketCapture
 import net.mullvad.mullvadvpn.test.e2e.router.packetCapture.PacketCaptureResult
@@ -96,8 +97,10 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
             val timestamp = System.currentTimeMillis()
             Attachment.saveAttachment("capture-testNegativeLeak-$timestamp.pcap", capturedPcap)
 
-            val leakRules = listOf(NoTrafficToHostRule(targetIpAddress))
-            LeakCheck.assertNoLeaks(capturedStreams, leakRules)
+            NetworkTrafficChecker.checkTrafficAgainstRules(
+                capturedStreams,
+                OnlyTrafficToHostRule(targetIpAddress)
+            )
         }
 
     @Test
@@ -156,8 +159,10 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
             val timestamp = System.currentTimeMillis()
             Attachment.saveAttachment("capture-testShouldHaveLeak-$timestamp.pcap", capturedPcap)
 
-            val leakRules = listOf(NoTrafficToHostRule(targetIpAddress))
-            LeakCheck.assertLeaks(capturedStreams, leakRules)
+            NetworkTrafficChecker.checkTrafficAgainstRules(
+                capturedStreams,
+                NoTrafficToHostRule(targetIpAddress)
+            )
         }
 
     @Test
@@ -212,8 +217,10 @@ class LeakTest : EndToEndTest(BuildConfig.FLAVOR_infrastructure) {
                 capturedPcap,
             )
 
-            val leakRules = listOf(NoTrafficToHostRule(targetIpAddress))
-            LeakCheck.assertLeaks(capturedStreams, leakRules)
+            NetworkTrafficChecker.checkTrafficAgainstRules(
+                capturedStreams,
+                OnlyTrafficToHostRule(targetIpAddress)
+            )
         }
 
     private fun disableObfuscation() {

android/test/e2e/src/main/kotlin/net/mullvad/mullvadvpn/test/e2e/misc/LeakCheck.kt

@@ -0,0 +1,36 @@
+package net.mullvad.mullvadvpn.test.e2e.misc
+
+import net.mullvad.mullvadvpn.test.e2e.router.packetCapture.Stream
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNotEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+
+object NetworkTrafficChecker {
+    fun checkTrafficAgainstRules(streams: List<Stream>, vararg rules: TrafficRule) {
+        // Assert that there are streams to be analyzed. Stream objects are guaranteed to contain
+        // packets when initialized.
+        assertTrue(streams.isNotEmpty(), "List of streams is empty.")
+
+        for (rule in rules) {
+            rule.assertTraffic(streams)
+        }
+    }
+}
+
+interface TrafficRule {
+    fun assertTraffic(streams: List<Stream>)
+}
+
+class OnlyTrafficToHostRule(private val host: String) : TrafficRule {
+    override fun assertTraffic(streams: List<Stream>) {
+        // TODO: Change to proper array/list assertion.
+        streams.forEach { assertEquals(host, it.destinationHost.ipAddress) }
+    }
+}
+
+class NoTrafficToHostRule(private val host: String) : TrafficRule {
+    override fun assertTraffic(streams: List<Stream>) {
+        // TODO: Change to proper array/list assertion.
+        streams.forEach { assertNotEquals(host, it.destinationHost.ipAddress) }
+    }
+}