ThreadStorageTest.java

/*
 * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */

package software.amazon.awssdk.threadcontext;

import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;

import static org.assertj.core.api.Assertions.assertThat;

class ThreadStorageTest {

    @AfterEach
    void cleanup() {
        ThreadStorage.clear();
    }

    @Test
    void putAndGet_shouldStoreAndRetrieveValue() {
        ThreadStorage.put("test-key", "test-value");
        
        assertThat(ThreadStorage.get("test-key")).isEqualTo("test-value");
    }

    @Test
    void get_withNonExistentKey_shouldReturnNull() {
        assertThat(ThreadStorage.get("non-existent")).isNull();
    }

    @Test
    void put_withValidKeyValue_shouldStoreValue() {
        ThreadStorage.put("test-key", "test-value");
        
        String removed = ThreadStorage.remove("test-key");
        
        assertThat(removed).isEqualTo("test-value");
        assertThat(ThreadStorage.get("test-key")).isNull();
    }

    @Test
    void remove_withExistingKey_shouldRemoveAndReturnValue() {
        ThreadStorage.put("test-key", "test-value");
        ThreadStorage.put("test-key", null);
        
        assertThat(ThreadStorage.get("test-key")).isNull();
    }

    @Test
    void clear_withMultipleValues_shouldRemoveAllValues() {
        ThreadStorage.put("key1", "value1");
        ThreadStorage.put("key2", "value2");
        
        ThreadStorage.clear();
        
        assertThat(ThreadStorage.get("key1")).isNull();
        assertThat(ThreadStorage.get("key2")).isNull();
    }
}