From efed30f903199e90be4f280c65f2f410b55835df Mon Sep 17 00:00:00 2001
From: abysssol <abysssol@pm.me>
Date: Mon, 11 Mar 2024 20:53:45 -0400
Subject: [PATCH] nixos/ollama: add test for the ollama service

---
 nixos/tests/all-tests.nix |  1 +
 nixos/tests/ollama.nix    | 56 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 nixos/tests/ollama.nix

diff --git a/nixos/tests/all-tests.nix b/nixos/tests/all-tests.nix
index 6c188593a97a..ac64b85dd486 100644
--- a/nixos/tests/all-tests.nix
+++ b/nixos/tests/all-tests.nix
@@ -640,6 +640,7 @@ in {
   nzbget = handleTest ./nzbget.nix {};
   nzbhydra2 = handleTest ./nzbhydra2.nix {};
   oh-my-zsh = handleTest ./oh-my-zsh.nix {};
+  ollama = handleTest ./ollama.nix {};
   ombi = handleTest ./ombi.nix {};
   openarena = handleTest ./openarena.nix {};
   openldap = handleTest ./openldap.nix {};
diff --git a/nixos/tests/ollama.nix b/nixos/tests/ollama.nix
new file mode 100644
index 000000000000..4b21f445cdbd
--- /dev/null
+++ b/nixos/tests/ollama.nix
@@ -0,0 +1,56 @@
+import ./make-test-python.nix ({ pkgs, lib, ... }:
+let
+  mainPort = "11434";
+  altPort = "11435";
+
+  curlRequest = port: request:
+    "curl http://127.0.0.1:${port}/api/generate -d '${builtins.toJSON request}'";
+
+  prompt = {
+    model = "tinydolphin";
+    prompt = "lorem ipsum";
+    options = {
+      seed = 69;
+      temperature = 0;
+    };
+  };
+in
+{
+  name = "ollama";
+  meta = with lib.maintainers; {
+    maintainers = [ abysssol ];
+  };
+
+  nodes = {
+    cpu = { ... }: {
+      services.ollama.enable = true;
+    };
+
+    rocm = { ... }: {
+      services.ollama.enable = true;
+      services.ollama.acceleration = "rocm";
+    };
+
+    cuda = { ... }: {
+      services.ollama.enable = true;
+      services.ollama.acceleration = "cuda";
+    };
+
+    altAddress = { ... }: {
+      services.ollama.enable = true;
+      services.ollama.listenAddress = "127.0.0.1:${altPort}";
+    };
+  };
+
+  testScript = ''
+    vms = [ cpu, rocm, cuda, altAddress ];
+
+    start_all()
+    for vm in vms:
+        vm.wait_for_unit("multi-user.target")
+
+    stdout = cpu.succeed("""${curlRequest mainPort prompt}""", timeout=100)
+
+    stdout = altAddress.succeed("""${curlRequest altPort prompt}""", timeout=100)
+  '';
+})