Test Context Management

/**
 * Thread-local SecurityContext management for testing environments
 * Provides isolation between test methods and threads
 */
public class TestSecurityContextHolder {
    
    /**
     * Set the SecurityContext for the current thread
     * @param context The SecurityContext to set
     */
    public static void setContext(SecurityContext context);
    
    /**
     * Get the current SecurityContext for this thread
     * @return The current SecurityContext or empty context if none set
     */
    public static SecurityContext getContext();
    
    /**
     * Clear the SecurityContext for the current thread
     * Should be called after each test to prevent context leakage
     */
    public static void clearContext();
    
    /**
     * Create and set SecurityContext with the provided Authentication
     * Convenience method that creates SecurityContext automatically
     * @param authentication The Authentication to set in a new SecurityContext
     */
    public static void setAuthentication(Authentication authentication);
}

@Test
public void testWithProgrammaticSecurityContext() {
    // Create custom authentication
    Authentication auth = new UsernamePasswordAuthenticationToken(
        "programmatic-user", 
        "password",
        Arrays.asList(new SimpleGrantedAuthority("ROLE_CUSTOM"))
    );
    
    // Set authentication in test context
    TestSecurityContextHolder.setAuthentication(auth);
    
    try {
        // Test code that requires security context
        SecurityContext context = TestSecurityContextHolder.getContext();
        assertThat(context.getAuthentication().getName()).isEqualTo("programmatic-user");
        
        // Test business logic that depends on security context
        String currentUser = getCurrentUsername(); // Your method
        assertThat(currentUser).isEqualTo("programmatic-user");
        
    } finally {
        // Always clean up
        TestSecurityContextHolder.clearContext();
    }
}

@Test
public void testWithCustomSecurityContext() {
    // Create custom SecurityContext
    SecurityContext customContext = SecurityContextHolder.createEmptyContext();
    
    // Create complex authentication with custom principal
    CustomUserPrincipal principal = new CustomUserPrincipal("custom-user", "department-A");
    Authentication auth = new UsernamePasswordAuthenticationToken(
        principal, 
        "credentials",
        Arrays.asList(
            new SimpleGrantedAuthority("ROLE_USER"),
            new SimpleGrantedAuthority("DEPT_A")
        )
    );
    
    customContext.setAuthentication(auth);
    TestSecurityContextHolder.setContext(customContext);
    
    try {
        // Test with custom context
        SecurityContext retrievedContext = TestSecurityContextHolder.getContext();
        assertThat(retrievedContext).isEqualTo(customContext);
        
        CustomUserPrincipal retrievedPrincipal = 
            (CustomUserPrincipal) retrievedContext.getAuthentication().getPrincipal();
        assertThat(retrievedPrincipal.getDepartment()).isEqualTo("department-A");
        
    } finally {
        TestSecurityContextHolder.clearContext();
    }
}

/**
 * Adapter that bridges TestSecurityContextHolder to SecurityContextHolderStrategy
 * Enables TestSecurityContextHolder to work with SecurityContextHolder strategy pattern
 */
public class TestSecurityContextHolderStrategyAdapter implements SecurityContextHolderStrategy {
    
    /**
     * Clear the SecurityContext using TestSecurityContextHolder
     */
    void clearContext();
    
    /**
     * Get SecurityContext using TestSecurityContextHolder
     * @return The current SecurityContext
     */
    SecurityContext getContext();
    
    /**
     * Set SecurityContext using TestSecurityContextHolder
     * @param context The SecurityContext to set
     */
    void setContext(SecurityContext context);
    
    /**
     * Create empty SecurityContext
     * @return New empty SecurityContext instance
     */
    SecurityContext createEmptyContext();
}

@Test
public void testWithStrategyAdapter() {
    TestSecurityContextHolderStrategyAdapter strategy = 
        new TestSecurityContextHolderStrategyAdapter();
    
    // Use strategy pattern
    SecurityContext context = strategy.createEmptyContext();
    Authentication auth = new UsernamePasswordAuthenticationToken(
        "strategy-user", "password", 
        Arrays.asList(new SimpleGrantedAuthority("ROLE_USER"))
    );
    context.setAuthentication(auth);
    
    strategy.setContext(context);
    
    try {
        SecurityContext retrievedContext = strategy.getContext();
        assertThat(retrievedContext.getAuthentication().getName())
            .isEqualTo("strategy-user");
    } finally {
        strategy.clearContext();
    }
}

/**
 * TestExecutionListener that processes @WithSecurityContext annotations
 * Handles security context setup and cleanup during test execution
 * Order: 10000 (runs after most Spring TestExecutionListeners)
 */
public class WithSecurityContextTestExecutionListener implements TestExecutionListener {
    
    /**
     * Process @WithSecurityContext annotations before test method execution
     * Sets up SecurityContext based on annotation configuration
     */
    void beforeTestMethod(TestContext testContext) throws Exception;
    
    /**
     * Process @WithSecurityContext annotations before test execution  
     * Alternative timing based on TestExecutionEvent configuration
     */
    void beforeTestExecution(TestContext testContext) throws Exception;
    
    /**
     * Clean up SecurityContext after test method completion
     * Ensures no context leakage between tests
     */
    void afterTestMethod(TestContext testContext) throws Exception;
    
    /**
     * Clean up SecurityContext after test execution
     * Alternative cleanup timing based on configuration
     */
    void afterTestExecution(TestContext testContext) throws Exception;
}

/**
 * TestExecutionListener for reactive security context management
 * Integrates SecurityContext with Reactor Context for WebFlux testing
 * Order: 11000 (runs after WithSecurityContextTestExecutionListener)
 */
public class ReactorContextTestExecutionListener implements TestExecutionListener {
    
    /**
     * Set up Reactor Context with SecurityContext before test execution
     * Enables reactive components to access security context
     */
    void beforeTestExecution(TestContext testContext) throws Exception;
    
    /**
     * Clean up Reactor Context after test execution
     * Prevents context pollution between tests
     */
    void afterTestExecution(TestContext testContext) throws Exception;
}

/**
 * TestExecutionListener that delegates to other listeners
 * Useful for conditional or dynamic listener management
 */
public class DelegatingTestExecutionListener implements TestExecutionListener {
    
    /**
     * Constructor with delegate listeners
     * @param delegates The TestExecutionListeners to delegate to
     */
    public DelegatingTestExecutionListener(TestExecutionListener... delegates);
    
    // Delegates all TestExecutionListener methods to configured delegates
}

/**
 * Utility class for accessing and manipulating Spring Security infrastructure
 * components during testing. Provides access to SecurityContextRepository,
 * CsrfTokenRepository, and related components from HttpServletRequest.
 */
public abstract class WebTestUtils {
    
    /**
     * Get SecurityContextRepository from request's filter chain
     * @param request The HttpServletRequest to examine
     * @return SecurityContextRepository or default HttpSessionSecurityContextRepository
     */
    public static SecurityContextRepository getSecurityContextRepository(HttpServletRequest request);
    
    /**
     * Set SecurityContextRepository for request's filter chain
     * @param request The HttpServletRequest to modify
     * @param securityContextRepository The SecurityContextRepository to set
     */
    public static void setSecurityContextRepository(
        HttpServletRequest request,
        SecurityContextRepository securityContextRepository
    );
    
    /**
     * Get CsrfTokenRepository from request's filter chain
     * @param request The HttpServletRequest to examine
     * @return CsrfTokenRepository or default HttpSessionCsrfTokenRepository
     */
    public static CsrfTokenRepository getCsrfTokenRepository(HttpServletRequest request);
    
    /**
     * Get CsrfTokenRequestHandler from request's filter chain
     * @param request The HttpServletRequest to examine
     * @return CsrfTokenRequestHandler or default XorCsrfTokenRequestAttributeHandler
     */
    public static CsrfTokenRequestHandler getCsrfTokenRequestHandler(HttpServletRequest request);
    
    /**
     * Set CsrfTokenRepository for request's filter chain
     * @param request The HttpServletRequest to modify
     * @param repository The CsrfTokenRepository to set
     */
    public static void setCsrfTokenRepository(
        HttpServletRequest request, 
        CsrfTokenRepository repository
    );
}

@SpringBootTest
@AutoConfigureTestDatabase
public class WebInfrastructureTest {
    
    @Autowired
    private WebApplicationContext context;
    
    private MockMvc mockMvc;
    
    @BeforeEach
    public void setup() {
        mockMvc = MockMvcBuilders
            .webAppContextSetup(context)
            .apply(springSecurity())
            .build();
    }
    
    @Test
    public void testCustomSecurityContextRepository() throws Exception {
        // Create custom SecurityContextRepository for testing
        SecurityContextRepository customRepo = new HttpSessionSecurityContextRepository();
        
        mockMvc.perform(get("/test-endpoint")
                .requestAttr("customRepo", customRepo)
                .with(user("testuser")))
               .andDo(result -> {
                   HttpServletRequest request = result.getRequest();
                   
                   // Get current repository
                   SecurityContextRepository currentRepo = 
                       WebTestUtils.getSecurityContextRepository(request);
                   assertThat(currentRepo).isNotNull();
                   
                   // Set custom repository
                   WebTestUtils.setSecurityContextRepository(request, customRepo);
                   
                   // Verify change
                   SecurityContextRepository updatedRepo = 
                       WebTestUtils.getSecurityContextRepository(request);
                   assertThat(updatedRepo).isEqualTo(customRepo);
               })
               .andExpect(status().isOk());
    }
    
    @Test
    public void testCsrfTokenRepositoryAccess() throws Exception {
        mockMvc.perform(post("/csrf-endpoint")
                .with(csrf())
                .with(user("testuser")))
               .andDo(result -> {
                   HttpServletRequest request = result.getRequest();
                   
                   // Access CSRF token repository
                   CsrfTokenRepository tokenRepo = 
                       WebTestUtils.getCsrfTokenRepository(request);
                   assertThat(tokenRepo).isNotNull();
                   
                   // Access CSRF token request handler
                   CsrfTokenRequestHandler requestHandler = 
                       WebTestUtils.getCsrfTokenRequestHandler(request);
                   assertThat(requestHandler).isNotNull();
               })
               .andExpect(status().isOk());
    }
    
    @Test
    public void testCustomCsrfConfiguration() throws Exception {
        // Create custom CSRF token repository
        CsrfTokenRepository customCsrfRepo = new HttpSessionCsrfTokenRepository();
        
        mockMvc.perform(post("/custom-csrf")
                .with(csrf())
                .with(user("testuser")))
               .andDo(result -> {
                   HttpServletRequest request = result.getRequest();
                   
                   // Set custom CSRF repository
                   WebTestUtils.setCsrfTokenRepository(request, customCsrfRepo);
                   
                   // Verify the change
                   CsrfTokenRepository retrievedRepo = 
                       WebTestUtils.getCsrfTokenRepository(request);
                   assertThat(retrievedRepo).isEqualTo(customCsrfRepo);
               })
               .andExpect(status().isOk());
    }
}

@SpringBootTest
public class ComplexSecurityContextTest {
    
    @Test
    public void testMultipleSecurityContexts() {
        // First context - regular user
        Authentication userAuth = new UsernamePasswordAuthenticationToken(
            "user", "password", 
            Arrays.asList(new SimpleGrantedAuthority("ROLE_USER"))
        );
        TestSecurityContextHolder.setAuthentication(userAuth);
        
        try {
            // Test user operations
            String userResult = performUserOperation();
            assertThat(userResult).contains("user");
            
            // Switch to admin context
            Authentication adminAuth = new UsernamePasswordAuthenticationToken(
                "admin", "password",
                Arrays.asList(
                    new SimpleGrantedAuthority("ROLE_ADMIN"),
                    new SimpleGrantedAuthority("ROLE_USER")
                )
            );
            TestSecurityContextHolder.setAuthentication(adminAuth);
            
            // Test admin operations
            String adminResult = performAdminOperation();
            assertThat(adminResult).contains("admin");
            
        } finally {
            TestSecurityContextHolder.clearContext();
        }
    }
    
    @Test
    public void testSecurityContextInCallbacks() {
        // Set up context for async operations
        Authentication auth = new UsernamePasswordAuthenticationToken(
            "async-user", "password",
            Arrays.asList(new SimpleGrantedAuthority("ROLE_USER"))
        );
        TestSecurityContextHolder.setAuthentication(auth);
        
        try {
            CompletableFuture<String> future = CompletableFuture.supplyAsync(() -> {
                // Context should be available in async callback
                SecurityContext context = TestSecurityContextHolder.getContext();
                return context.getAuthentication().getName();
            });
            
            String result = future.get(1, TimeUnit.SECONDS);
            assertThat(result).isEqualTo("async-user");
            
        } catch (Exception e) {
            fail("Async operation failed", e);
        } finally {
            TestSecurityContextHolder.clearContext();
        }
    }
}

/**
 * Custom TestExecutionListener for specialized security context management
 */
public class CustomSecurityTestExecutionListener implements TestExecutionListener {
    
    @Override
    public void beforeTestMethod(TestContext testContext) throws Exception {
        Method testMethod = testContext.getTestMethod();
        
        // Check for custom security annotation
        CustomSecurity customSecurity = testMethod.getAnnotation(CustomSecurity.class);
        if (customSecurity != null) {
            setupCustomSecurityContext(customSecurity);
        }
    }
    
    @Override
    public void afterTestMethod(TestContext testContext) throws Exception {
        // Always clean up
        TestSecurityContextHolder.clearContext();
    }
    
    private void setupCustomSecurityContext(CustomSecurity annotation) {
        // Create custom authentication based on annotation
        Authentication auth = createCustomAuthentication(
            annotation.username(), 
            annotation.permissions()
        );
        TestSecurityContextHolder.setAuthentication(auth);
    }
    
    private Authentication createCustomAuthentication(String username, String[] permissions) {
        List<GrantedAuthority> authorities = Arrays.stream(permissions)
            .map(SimpleGrantedAuthority::new)
            .collect(Collectors.toList());
            
        return new UsernamePasswordAuthenticationToken(username, "password", authorities);
    }
}

// Usage with custom annotation
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface CustomSecurity {
    String username() default "testuser";
    String[] permissions() default {};
}

// Test class using custom listener
@SpringBootTest
@TestExecutionListeners(
    listeners = CustomSecurityTestExecutionListener.class,
    mergeMode = TestExecutionListeners.MergeMode.MERGE_WITH_DEFAULTS
)
public class CustomSecurityTest {
    
    @Test
    @CustomSecurity(username = "custom-user", permissions = {"READ", "WRITE"})
    public void testWithCustomSecurity() {
        SecurityContext context = TestSecurityContextHolder.getContext();
        assertThat(context.getAuthentication().getName()).isEqualTo("custom-user");
        
        Collection<? extends GrantedAuthority> authorities = 
            context.getAuthentication().getAuthorities();
        assertThat(authorities).hasSize(2);
        assertThat(authorities.stream().map(GrantedAuthority::getAuthority))
            .containsExactlyInAnyOrder("READ", "WRITE");
    }
}

@SpringBootTest  
public class IntegratedContextManagementTest {
    
    @Autowired
    private WebApplicationContext context;
    
    private MockMvc mockMvc;
    
    @BeforeEach
    public void setup() {
        mockMvc = MockMvcBuilders
            .webAppContextSetup(context)
            .apply(springSecurity())
            .build();
    }
    
    @Test
    public void testMixedContextUsage() throws Exception {
        // Set up context programmatically
        Authentication auth = new UsernamePasswordAuthenticationToken(
            "mixed-user", "password",
            Arrays.asList(new SimpleGrantedAuthority("ROLE_USER"))
        );
        TestSecurityContextHolder.setAuthentication(auth);
        
        try {
            // Use with MockMvc via testSecurityContext()
            mockMvc.perform(get("/user-info")
                    .with(testSecurityContext()))
                   .andExpect(status().isOk())
                   .andExpect(authenticated().withUsername("mixed-user"));
                   
            // Verify context is still available
            SecurityContext retrievedContext = TestSecurityContextHolder.getContext();
            assertThat(retrievedContext.getAuthentication().getName())
                .isEqualTo("mixed-user");
                
        } finally {
            TestSecurityContextHolder.clearContext();
        }
    }
}

@SpringBootTest
public class ThreadSafetyTest {
    
    @Test
    public void testContextIsolationBetweenThreads() throws Exception {
        // Main thread context
        TestSecurityContextHolder.setAuthentication(
            new UsernamePasswordAuthenticationToken("main-user", "password", 
                Arrays.asList(new SimpleGrantedAuthority("ROLE_USER")))
        );
        
        try {
            // Verify main thread context
            assertThat(TestSecurityContextHolder.getContext().getAuthentication().getName())
                .isEqualTo("main-user");
            
            // Test in separate thread
            CompletableFuture<String> otherThreadResult = CompletableFuture.supplyAsync(() -> {
                // Other thread should have empty context
                SecurityContext otherContext = TestSecurityContextHolder.getContext();
                return otherContext.getAuthentication() != null ? 
                    otherContext.getAuthentication().getName() : "no-auth";
            });
            
            String result = otherThreadResult.get(1, TimeUnit.SECONDS);
            assertThat(result).isEqualTo("no-auth");
            
            // Main thread context should still be intact
            assertThat(TestSecurityContextHolder.getContext().getAuthentication().getName())
                .isEqualTo("main-user");
                
        } finally {
            TestSecurityContextHolder.clearContext();
        }
    }
}