sqrtspace-dotnet/tests/SqrtSpace.SpaceTime.Tests/EntityFramework/SpaceTimeDbContextTests.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using SqrtSpace.SpaceTime.EntityFramework;
using Xunit;

namespace SqrtSpace.SpaceTime.Tests.EntityFramework;

public class SpaceTimeDbContextTests : IDisposable
{
    private readonly ServiceProvider _serviceProvider;
    private readonly TestDbContext _context;

    public SpaceTimeDbContextTests()
    {
        var services = new ServiceCollection();
        services.AddDbContext<TestDbContext>(options =>
        {
            options.UseInMemoryDatabase($"TestDb_{Guid.NewGuid()}");
            options.UseSpaceTimeOptimizer(opt =>
            {
                opt.EnableSqrtNChangeTracking = true;
                opt.BufferPoolStrategy = BufferPoolStrategy.SqrtN;
                opt.EnableQueryCheckpointing = true;
                opt.MaxTrackedEntities = 100;
            });
        });

        _serviceProvider = services.BuildServiceProvider();
        _context = _serviceProvider.GetRequiredService<TestDbContext>();
        _context.Database.EnsureCreated();
    }

    public void Dispose()
    {
        _context?.Dispose();
        _serviceProvider?.Dispose();
    }

    [Fact]
    public async Task ToListWithSqrtNMemoryAsync_ReturnsAllEntities()
    {
        // Arrange
        var customers = GenerateCustomers(100);
        _context.Customers.AddRange(customers);
        await _context.SaveChangesAsync();

        // Act
        var result = await _context.Customers
            .Where(c => c.IsActive)
            .ToListWithSqrtNMemoryAsync();

        // Assert
        result.Should().HaveCount(50); // Half are active
        result.All(c => c.IsActive).Should().BeTrue();
    }

    [Fact]
    public async Task BatchBySqrtNAsync_ProcessesInBatches()
    {
        // Arrange
        var orders = GenerateOrders(1000);
        _context.Orders.AddRange(orders);
        await _context.SaveChangesAsync();

        // Act
        var batchCount = 0;
        var totalProcessed = 0;
        
        await foreach (var batch in _context.Orders.BatchBySqrtNAsync())
        {
            batchCount++;
            totalProcessed += batch.Count();
        }

        // Assert
        batchCount.Should().BeGreaterThan(1);
        totalProcessed.Should().Be(1000);
        var expectedBatchSize = (int)Math.Sqrt(1000); // ~31
        batchCount.Should().BeCloseTo(1000 / expectedBatchSize, 2);
    }

    [Fact]
    public async Task SqrtNChangeTracking_LimitsTrackedEntities()
    {
        // Arrange
        var customers = GenerateCustomers(200);
        _context.Customers.AddRange(customers);
        await _context.SaveChangesAsync();
        _context.ChangeTracker.Clear();

        // Act
        // Load entities in batches, change tracking should limit memory
        var loaded = new List<Customer>();
        await foreach (var batch in _context.Customers.BatchBySqrtNAsync())
        {
            loaded.AddRange(batch);
            
            // Modify some entities
            foreach (var customer in batch.Take(5))
            {
                customer.Name += " Modified";
            }
        }

        // Assert
        loaded.Should().HaveCount(200);
        // Change tracker should have limited entries due to √n tracking
        var trackedCount = _context.ChangeTracker.Entries<Customer>().Count();
        trackedCount.Should().BeLessThanOrEqualTo(100); // MaxTrackedEntities
    }

    [Fact]
    public async Task QueryCheckpointing_EnablesRecovery()
    {
        // Arrange
        var products = GenerateProducts(500);
        _context.Products.AddRange(products);
        await _context.SaveChangesAsync();

        var checkpointId = Guid.NewGuid().ToString();
        var processedIds = new List<int>();

        // Act - Simulate partial processing
        try
        {
            await foreach (var batch in _context.Products
                .OrderBy(p => p.Id)
                .BatchBySqrtNAsync(checkpointId))
            {
                foreach (var product in batch)
                {
                    processedIds.Add(product.Id);
                    
                    // Simulate failure after processing 100 items
                    if (processedIds.Count == 100)
                    {
                        throw new Exception("Simulated failure");
                    }
                }
            }
        }
        catch
        {
            // Expected failure
        }

        // Resume from checkpoint
        var resumedIds = new List<int>();
        await foreach (var batch in _context.Products
            .OrderBy(p => p.Id)
            .BatchBySqrtNAsync(checkpointId, resumeFromCheckpoint: true))
        {
            foreach (var product in batch)
            {
                resumedIds.Add(product.Id);
            }
        }

        // Assert
        processedIds.Should().HaveCount(100);
        resumedIds.Should().HaveCountGreaterThan(300); // Should continue from checkpoint
        resumedIds.Should().NotContain(processedIds); // Should not reprocess
    }

    [Fact]
    public async Task ExternalSortingQuery_HandlesLargeDataset()
    {
        // Arrange
        var orders = GenerateOrders(1000);
        _context.Orders.AddRange(orders);
        await _context.SaveChangesAsync();

        // Act
        var sorted = await _context.Orders
            .UseExternalSorting()
            .OrderBy(o => o.OrderDate)
            .ThenBy(o => o.TotalAmount)
            .ToListAsync();

        // Assert
        sorted.Should().HaveCount(1000);
        sorted.Should().BeInAscendingOrder(o => o.OrderDate)
            .And.ThenBeInAscendingOrder(o => o.TotalAmount);
    }

    [Fact]
    public async Task StreamQueryResultsAsync_StreamsEfficiently()
    {
        // Arrange
        var customers = GenerateCustomers(500);
        _context.Customers.AddRange(customers);
        await _context.SaveChangesAsync();

        // Act
        var streamed = new List<Customer>();
        await foreach (var customer in _context.Customers
            .Where(c => c.CreatedDate > DateTime.Today.AddDays(-30))
            .StreamQueryResultsAsync())
        {
            streamed.Add(customer);
        }

        // Assert
        streamed.Should().HaveCountGreaterThan(0);
        streamed.All(c => c.CreatedDate > DateTime.Today.AddDays(-30)).Should().BeTrue();
    }

    [Fact]
    public async Task BufferPoolStrategy_OptimizesMemoryUsage()
    {
        // Arrange
        var orders = GenerateOrders(10000);
        
        // Act - Use buffered save changes
        await _context.BulkInsertWithSqrtNBufferingAsync(orders);

        // Assert
        var count = await _context.Orders.CountAsync();
        count.Should().Be(10000);
    }

    [Fact]
    public async Task ComplexQuery_WithSpaceTimeOptimizations()
    {
        // Arrange
        await SeedComplexDataAsync();

        // Act
        var result = await _context.Orders
            .Include(o => o.OrderItems)
            .ThenInclude(oi => oi.Product)
            .Where(o => o.OrderDate >= DateTime.Today.AddMonths(-1))
            .UseExternalSorting()
            .GroupBy(o => o.Customer.City)
            .Select(g => new
            {
                City = g.Key,
                OrderCount = g.Count(),
                TotalRevenue = g.Sum(o => o.TotalAmount),
                AverageOrderValue = g.Average(o => o.TotalAmount)
            })
            .OrderByDescending(x => x.TotalRevenue)
            .ToListWithSqrtNMemoryAsync();

        // Assert
        result.Should().NotBeEmpty();
        result.Should().BeInDescendingOrder(x => x.TotalRevenue);
    }

    [Fact]
    public async Task ChangeTracking_WithAutoDetectChanges()
    {
        // Arrange
        var customer = new Customer { Name = "Test Customer", IsActive = true };
        _context.Customers.Add(customer);
        await _context.SaveChangesAsync();

        // Act
        customer.Name = "Updated Customer";
        customer.Email = "updated@example.com";

        // Assert
        var entry = _context.Entry(customer);
        entry.State.Should().Be(EntityState.Modified);
        var modifiedProps = entry.Properties
            .Where(p => p.IsModified)
            .Select(p => p.Metadata.Name)
            .ToList();
        modifiedProps.Should().Contain(new[] { "Name", "Email" });
    }

    [Fact]
    public async Task TransactionWithCheckpointing_MaintainsConsistency()
    {
        // Arrange
        var orders = GenerateOrders(100);

        // Act
        using var transaction = await _context.Database.BeginTransactionAsync();
        try
        {
            await _context.BulkInsertWithSqrtNBufferingAsync(orders);
            
            // Simulate some updates
            var ordersToUpdate = await _context.Orders
                .Take(10)
                .ToListAsync();
            
            foreach (var order in ordersToUpdate)
            {
                order.Status = "Processed";
            }
            
            await _context.SaveChangesAsync();
            await transaction.CommitAsync();
        }
        catch
        {
            await transaction.RollbackAsync();
            throw;
        }

        // Assert
        var processedCount = await _context.Orders
            .CountAsync(o => o.Status == "Processed");
        processedCount.Should().Be(10);
    }

    private List<Customer> GenerateCustomers(int count)
    {
        return Enumerable.Range(1, count).Select(i => new Customer
        {
            Name = $"Customer {i}",
            Email = $"customer{i}@example.com",
            City = $"City{i % 10}",
            IsActive = i % 2 == 0,
            CreatedDate = DateTime.Today.AddDays(-Random.Shared.Next(365))
        }).ToList();
    }

    private List<Order> GenerateOrders(int count)
    {
        return Enumerable.Range(1, count).Select(i => new Order
        {
            OrderNumber = $"ORD{i:D6}",
            OrderDate = DateTime.Today.AddDays(-Random.Shared.Next(365)),
            TotalAmount = Random.Shared.Next(10, 1000),
            Status = "Pending",
            CustomerId = Random.Shared.Next(1, 100)
        }).ToList();
    }

    private List<Product> GenerateProducts(int count)
    {
        return Enumerable.Range(1, count).Select(i => new Product
        {
            Name = $"Product {i}",
            SKU = $"SKU{i:D6}",
            Price = Random.Shared.Next(10, 500),
            StockQuantity = Random.Shared.Next(0, 100)
        }).ToList();
    }

    private async Task SeedComplexDataAsync()
    {
        var customers = GenerateCustomers(100);
        var products = GenerateProducts(50);
        
        _context.Customers.AddRange(customers);
        _context.Products.AddRange(products);
        await _context.SaveChangesAsync();

        var orders = new List<Order>();
        foreach (var customer in customers.Take(50))
        {
            for (int i = 0; i < Random.Shared.Next(1, 5); i++)
            {
                var order = new Order
                {
                    OrderNumber = $"ORD{Guid.NewGuid():N}",
                    OrderDate = DateTime.Today.AddDays(-Random.Shared.Next(60)),
                    CustomerId = customer.Id,
                    Status = "Pending",
                    OrderItems = new List<OrderItem>()
                };

                var itemCount = Random.Shared.Next(1, 5);
                for (int j = 0; j < itemCount; j++)
                {
                    var product = products[Random.Shared.Next(products.Count)];
                    var quantity = Random.Shared.Next(1, 10);
                    order.OrderItems.Add(new OrderItem
                    {
                        ProductId = product.Id,
                        Quantity = quantity,
                        UnitPrice = product.Price,
                        TotalPrice = product.Price * quantity
                    });
                }

                order.TotalAmount = order.OrderItems.Sum(oi => oi.TotalPrice);
                orders.Add(order);
            }
        }

        _context.Orders.AddRange(orders);
        await _context.SaveChangesAsync();
    }
}

public class TestDbContext : DbContext
{
    public TestDbContext(DbContextOptions<TestDbContext> options) : base(options) { }

    public DbSet<Customer> Customers { get; set; }
    public DbSet<Order> Orders { get; set; }
    public DbSet<Product> Products { get; set; }
    public DbSet<OrderItem> OrderItems { get; set; }

    protected override void OnModelCreating(ModelBuilder modelBuilder)
    {
        modelBuilder.Entity<Order>()
            .HasOne(o => o.Customer)
            .WithMany(c => c.Orders)
            .HasForeignKey(o => o.CustomerId);

        modelBuilder.Entity<OrderItem>()
            .HasOne(oi => oi.Order)
            .WithMany(o => o.OrderItems)
            .HasForeignKey(oi => oi.OrderId);

        modelBuilder.Entity<OrderItem>()
            .HasOne(oi => oi.Product)
            .WithMany()
            .HasForeignKey(oi => oi.ProductId);
    }
}

public class Customer
{
    public int Id { get; set; }
    public string Name { get; set; } = "";
    public string? Email { get; set; }
    public string City { get; set; } = "";
    public bool IsActive { get; set; }
    public DateTime CreatedDate { get; set; }
    public List<Order> Orders { get; set; } = new();
}

public class Order
{
    public int Id { get; set; }
    public string OrderNumber { get; set; } = "";
    public DateTime OrderDate { get; set; }
    public decimal TotalAmount { get; set; }
    public string Status { get; set; } = "";
    public int CustomerId { get; set; }
    public Customer Customer { get; set; } = null!;
    public List<OrderItem> OrderItems { get; set; } = new();
}

public class Product
{
    public int Id { get; set; }
    public string Name { get; set; } = "";
    public string SKU { get; set; } = "";
    public decimal Price { get; set; }
    public int StockQuantity { get; set; }
}

public class OrderItem
{
    public int Id { get; set; }
    public int OrderId { get; set; }
    public Order Order { get; set; } = null!;
    public int ProductId { get; set; }
    public Product Product { get; set; } = null!;
    public int Quantity { get; set; }
    public decimal UnitPrice { get; set; }
    public decimal TotalPrice { get; set; }
}