1 files changed, 56 insertions, 0 deletions
diff --git a/nixpkgs/nixos/tests/ollama.nix b/nixpkgs/nixos/tests/ollama.nix
new file mode 100644
index 000000000000..4b21f445cdbd
--- /dev/null
+++ b/nixpkgs/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)
+  '';
+})