From cf5906a82689ae0d6772a10ce1818094f1fb9b1c Mon Sep 17 00:00:00 2001
From: Virgil <virgil@lethean.io>
Date: Wed, 1 Apr 2026 09:27:55 +0000
Subject: [PATCH] feat(ansible): add group_by module support

Co-Authored-By: Virgil <virgil@lethean.io>
---
 mock_ssh_test.go    |  2 ++
 modules.go          | 55 +++++++++++++++++++++++++++++++++++++++++++++
 modules_adv_test.go | 42 ++++++++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/mock_ssh_test.go b/mock_ssh_test.go
index 6100282..ee0c106 100644
--- a/mock_ssh_test.go
+++ b/mock_ssh_test.go
@@ -522,6 +522,8 @@ func executeModuleWithMock(e *Executor, mock *MockSSHClient, host string, task *
 	// Inventory mutation
 	case "ansible.builtin.add_host":
 		return e.moduleAddHost(args)
+	case "ansible.builtin.group_by":
+		return e.moduleGroupBy(args, host)
 
 	// SSH keys
 	case "ansible.posix.authorized_key", "ansible.builtin.authorized_key":
diff --git a/modules.go b/modules.go
index 60c7e7d..b4585de 100644
--- a/modules.go
+++ b/modules.go
@@ -131,6 +131,8 @@ func (e *Executor) executeModule(ctx context.Context, host string, client *SSHCl
 		return e.moduleIncludeVars(args)
 	case "ansible.builtin.add_host":
 		return e.moduleAddHost(args)
+	case "ansible.builtin.group_by":
+		return e.moduleGroupBy(args, host)
 	case "ansible.builtin.meta":
 		return e.moduleMeta(args)
 	case "ansible.builtin.setup":
@@ -1474,6 +1476,59 @@ func (e *Executor) moduleAddHost(args map[string]any) (*TaskResult, error) {
 	return &TaskResult{Changed: true, Msg: "host added: " + name}, nil
 }
 
+func (e *Executor) moduleGroupBy(args map[string]any, host string) (*TaskResult, error) {
+	groupName := getStringArg(args, "key", "")
+	if groupName == "" {
+		groupName = getStringArg(args, "name", "")
+	}
+	if groupName == "" {
+		groupName = getStringArg(args, "_raw_params", "")
+	}
+	if groupName == "" {
+		return nil, coreerr.E("Executor.moduleGroupBy", "key required", nil)
+	}
+	if host == "" {
+		host = getStringArg(args, "inventory_hostname", "")
+	}
+	if host == "" {
+		host = getStringArg(args, "host", "")
+	}
+	if host == "" {
+		return nil, coreerr.E("Executor.moduleGroupBy", "host required", nil)
+	}
+
+	if e.inventory == nil {
+		e.inventory = &Inventory{}
+	}
+	if e.inventory.All == nil {
+		e.inventory.All = &InventoryGroup{}
+	}
+	if e.inventory.All.Children == nil {
+		e.inventory.All.Children = make(map[string]*InventoryGroup)
+	}
+	if e.inventory.All.Hosts == nil {
+		e.inventory.All.Hosts = make(map[string]*Host)
+	}
+
+	hostEntry := e.inventory.All.Hosts[host]
+	if hostEntry == nil {
+		hostEntry = &Host{}
+		e.inventory.All.Hosts[host] = hostEntry
+	}
+
+	group := e.inventory.All.Children[groupName]
+	if group == nil {
+		group = &InventoryGroup{}
+		e.inventory.All.Children[groupName] = group
+	}
+	if group.Hosts == nil {
+		group.Hosts = make(map[string]*Host)
+	}
+	group.Hosts[host] = hostEntry
+
+	return &TaskResult{Changed: true, Msg: "host grouped: " + host + " -> " + groupName}, nil
+}
+
 func (e *Executor) moduleMeta(args map[string]any) (*TaskResult, error) {
 	// meta module controls play execution
 	// Most actions are no-ops for us
diff --git a/modules_adv_test.go b/modules_adv_test.go
index 60f5228..0be87ff 100644
--- a/modules_adv_test.go
+++ b/modules_adv_test.go
@@ -1103,6 +1103,30 @@ func TestModulesAdv_ModuleAddHost_Good_AddsHostAndGroups(t *testing.T) {
 	assert.Equal(t, "app", GetHostVars(e.inventory, "web2")["role"])
 }
 
+func TestModulesAdv_ModuleGroupBy_Good_AddsHostToDerivedGroup(t *testing.T) {
+	e := NewExecutor("/tmp")
+	e.SetInventoryDirect(&Inventory{
+		All: &InventoryGroup{
+			Hosts: map[string]*Host{
+				"web1": {AnsibleHost: "10.0.0.1"},
+			},
+			Children: map[string]*InventoryGroup{},
+		},
+	})
+
+	result, err := e.moduleGroupBy(map[string]any{
+		"key": "role_web",
+	}, "web1")
+
+	require.NoError(t, err)
+	assert.True(t, result.Changed)
+	assert.Equal(t, "host grouped: web1 -> role_web", result.Msg)
+	require.NotNil(t, e.inventory.All.Children["role_web"])
+	assert.Contains(t, e.inventory.All.Children["role_web"].Hosts, "web1")
+	assert.Equal(t, e.inventory.All.Hosts["web1"], e.inventory.All.Children["role_web"].Hosts["web1"])
+	assert.Equal(t, []string{"web1"}, GetHosts(e.inventory, "role_web"))
+}
+
 // --- Cross-module dispatch tests for advanced modules ---
 
 func TestModulesAdv_ExecuteModuleWithMock_Good_DispatchUser(t *testing.T) {
@@ -1140,6 +1164,24 @@ func TestModulesAdv_ExecuteModuleWithMock_Good_DispatchGroup(t *testing.T) {
 	assert.True(t, result.Changed)
 }
 
+func TestModulesAdv_ExecuteModuleWithMock_Good_DispatchGroupBy(t *testing.T) {
+	e, _ := newTestExecutorWithMock("host1")
+
+	task := &Task{
+		Module: "group_by",
+		Args: map[string]any{
+			"key": "role_app",
+		},
+	}
+
+	result, err := executeModuleWithMock(e, nil, "host1", task)
+
+	require.NoError(t, err)
+	assert.True(t, result.Changed)
+	assert.Contains(t, e.inventory.All.Children, "role_app")
+	assert.Contains(t, e.inventory.All.Children["role_app"].Hosts, "host1")
+}
+
 func TestModulesAdv_ExecuteModuleWithMock_Good_DispatchCron(t *testing.T) {
 	e, mock := newTestExecutorWithMock("host1")
 	mock.expectCommand(`crontab`, "", "", 0)