From fb561700856a844a6c076f07d9e70642b4b1f3e2 Mon Sep 17 00:00:00 2001
From: Szymon Sasin <szymon.sasin@gmail.com>
Date: Wed, 10 Jul 2024 19:49:40 +0300
Subject: [PATCH] fix

---
 .../mbed/coap/transport/TransportContext.java | 35 +++++++------------
 .../coap/transport/TransportContextTest.java  | 25 +++----------
 2 files changed, 17 insertions(+), 43 deletions(-)

diff --git a/coap-core/src/main/java/com/mbed/coap/transport/TransportContext.java b/coap-core/src/main/java/com/mbed/coap/transport/TransportContext.java
index ebf70ff0..53f26608 100644
--- a/coap-core/src/main/java/com/mbed/coap/transport/TransportContext.java
+++ b/coap-core/src/main/java/com/mbed/coap/transport/TransportContext.java
@@ -17,11 +17,11 @@
 
 import static java.util.Objects.requireNonNull;
 import java.time.Duration;
-import java.util.Iterator;
-import java.util.NoSuchElementException;
+import java.util.HashSet;
 import java.util.Objects;
+import java.util.Set;
 
-public final class TransportContext implements Iterable<TransportContext.Key<?>> {
+public final class TransportContext {
 
     private final Key key;
     private final Object value;
@@ -96,26 +96,17 @@ public int hashCode() {
         return Objects.hash(key, value, next);
     }
 
-    @Override
-    public Iterator<Key<?>> iterator() {
-        return new Iterator<Key<?>>() {
-            private TransportContext current = TransportContext.this;
-
-            @Override
-            public boolean hasNext() {
-                return current != null;
-            }
+    public Set<Key<?>> keys() {
+        Set<Key<?>> keys = new HashSet<>();
+        addKey(keys);
+        return keys;
+    }
 
-            @Override
-            public Key<?> next() {
-                if (!hasNext()) {
-                    throw new NoSuchElementException();
-                }
-                Key<?> key = current.key;
-                current = current.next;
-                return key;
-            }
-        };
+    private void addKey(Set<Key<?>> keys) {
+        keys.add(key);
+        if (next != null) {
+            next.addKey(keys);
+        }
     }
 
     public static final class Key<T> {
diff --git a/coap-core/src/test/java/com/mbed/coap/transport/TransportContextTest.java b/coap-core/src/test/java/com/mbed/coap/transport/TransportContextTest.java
index 322807ee..5162eade 100644
--- a/coap-core/src/test/java/com/mbed/coap/transport/TransportContextTest.java
+++ b/coap-core/src/test/java/com/mbed/coap/transport/TransportContextTest.java
@@ -16,13 +16,9 @@
 package com.mbed.coap.transport;
 
 import static com.mbed.coap.transport.TransportContext.EMPTY;
+import static org.assertj.core.util.Sets.set;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import java.util.Iterator;
-import java.util.NoSuchElementException;
 import nl.jqno.equalsverifier.EqualsVerifier;
 import nl.jqno.equalsverifier.Warning;
 import org.junit.jupiter.api.Test;
@@ -53,6 +49,7 @@ void merge() {
 
         assertEquals("111", ctx3.get(DUMMY_KEY));
         assertEquals("222", ctx3.get(DUMMY_KEY2));
+        assertEquals(set(DUMMY_KEY, DUMMY_KEY2), ctx3.keys());
     }
 
     @Test
@@ -61,6 +58,7 @@ void mergeWithEmpty() {
 
         assertEquals(ctx1, ctx1.with(EMPTY));
         assertEquals(ctx1, EMPTY.with(ctx1));
+        assertEquals(set(DUMMY_KEY), ctx1.keys());
     }
 
     @Test
@@ -72,22 +70,7 @@ void mergeAndOverWrite() {
 
         assertEquals("aaa", ctx3.get(DUMMY_KEY));
         assertEquals("222", ctx3.get(DUMMY_KEY2));
-    }
-
-    @Test
-    void listKeys() {
-        TransportContext ctx1 = TransportContext.of(DUMMY_KEY, "111").with(DUMMY_KEY2, "222");
-
-        Iterator<TransportContext.Key<?>> iterator = ctx1.iterator();
-
-        assertTrue(iterator.hasNext());
-        assertEquals(DUMMY_KEY2, iterator.next());
-
-        assertTrue(iterator.hasNext());
-        assertEquals(DUMMY_KEY, iterator.next());
-
-        assertFalse(iterator.hasNext());
-        assertThrows(NoSuchElementException.class, iterator::next);
+        assertEquals(set(DUMMY_KEY, DUMMY_KEY2), ctx3.keys());
     }
 
     @Test