diff --git a/LICENSE b/LICENSE
index 497106c..c63d87e 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,4 +1,4 @@
-Copyright © 2020, UChicago Argonne, LLC
+Copyright © 2020-2021, UChicago Argonne, LLC
 
 All Rights Reserved
 
diff --git a/Makefile b/Makefile
deleted file mode 100644
index 6fa28d6..0000000
--- a/Makefile
+++ /dev/null
@@ -1,13 +0,0 @@
-JULIA := julia --color=yes
-JULIA_SYSIMAGE_ARGS := $(JULIA_ARGS) --sysimage build/sysimage.so
-
-all: test
-
-build/sysimage.so: Manifest.toml Project.toml
-	mkdir -p build
-	$(JULIA) --project=test test/sysimage.jl
-
-test: build/sysimage.so
-	$(JULIA) --sysimage build/sysimage.so --project=test test/runtests.jl
-
-.PHONY: test test-python test-julia test-watch docs install
diff --git a/Manifest.toml b/Manifest.toml
index fa6f596..15f7a1c 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -1,10 +1,10 @@
 # This file is machine-generated - editing it directly is not advised
 
+[[ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+
 [[Artifacts]]
-deps = ["Pkg"]
-git-tree-sha1 = "c30985d8821e0cd73870b17b0ed0ce6dc44cb744"
 uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
-version = "1.3.0"
 
 [[Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -53,15 +53,13 @@ version = "0.3.0"
 
 [[Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
-git-tree-sha1 = "4fecfd5485d3c5de4003e19f00c6898cccd40667"
+git-tree-sha1 = "ac4132ad78082518ec2037ae5770b6e796f7f956"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.26.0"
+version = "3.27.0"
 
 [[CompilerSupportLibraries_jll]]
-deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "8e695f735fca77e9708e795eda62afdb869cbb70"
+deps = ["Artifacts", "Libdl"]
 uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
-version = "0.3.4+0"
 
 [[Conda]]
 deps = ["JSON", "VersionParsing"]
@@ -99,6 +97,10 @@ version = "1.0.2"
 deps = ["Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
+[[Downloads]]
+deps = ["ArgTools", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+
 [[ForwardDiff]]
 deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "NaNMath", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
 git-tree-sha1 = "e2af66012e08966366a43251e1fd421522908be6"
@@ -132,12 +134,6 @@ git-tree-sha1 = "81690084b6198a2e1da36fcfda16eeca9f9f24e4"
 uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 version = "0.21.1"
 
-[[JSON2]]
-deps = ["Dates", "Parsers", "Test"]
-git-tree-sha1 = "66397cc6c08922f98a28ab05a8d3002f9853b129"
-uuid = "2535ab7d-5cd8-5a07-80ac-9b1792aadce3"
-version = "0.3.2"
-
 [[JSONSchema]]
 deps = ["HTTP", "JSON", "ZipFile"]
 git-tree-sha1 = "b84ab8139afde82c7c65ba2b792fe12e01dd7307"
@@ -150,10 +146,22 @@ git-tree-sha1 = "e952f49e2242fa21edcf27bbd6c67041685bee5d"
 uuid = "4076af6c-e467-56ae-b986-b466b2749572"
 version = "0.21.6"
 
+[[LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+
+[[LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+
 [[LibGit2]]
-deps = ["Printf"]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
+[[LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+
 [[Libdl]]
 uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 
@@ -187,19 +195,20 @@ uuid = "739be429-bea8-5141-9913-cc70e7f3736d"
 version = "1.0.3"
 
 [[MbedTLS_jll]]
-deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "0eef589dd1c26a3ac9d753fe1a8bcad63f956fa6"
+deps = ["Artifacts", "Libdl"]
 uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
-version = "2.16.8+1"
 
 [[Mmap]]
 uuid = "a63ad114-7e13-5084-954f-fe012c677804"
 
+[[MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+
 [[MutableArithmetics]]
 deps = ["LinearAlgebra", "SparseArrays", "Test"]
-git-tree-sha1 = "6b6bb8f550dc38310afd4a0af0786dc3222459e2"
+git-tree-sha1 = "ff3aa3e4dbc837f80c2031de2f90125c8b3793f3"
 uuid = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
-version = "0.2.14"
+version = "0.2.15"
 
 [[NaNMath]]
 git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
@@ -207,9 +216,7 @@ uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 version = "0.3.5"
 
 [[NetworkOptions]]
-git-tree-sha1 = "ed3157f48a05543cce9b241e1f2815f7e843d96e"
 uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
-version = "1.2.0"
 
 [[OpenSpecFun_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
@@ -222,12 +229,6 @@ git-tree-sha1 = "4fa2ba51070ec13fcc7517db714445b4ab986bdf"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 version = "1.4.0"
 
-[[PackageCompiler]]
-deps = ["Libdl", "Pkg", "UUIDs"]
-git-tree-sha1 = "d448727c4b86be81b225b738c88d30334fda6779"
-uuid = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
-version = "1.2.5"
-
 [[Parsers]]
 deps = ["Dates"]
 git-tree-sha1 = "c8abc88faa3f7a3950832ac5d6e690881590d6dc"
@@ -235,7 +236,7 @@ uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
 version = "1.1.0"
 
 [[Pkg]]
-deps = ["Dates", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "UUIDs"]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [[Printf]]
@@ -249,7 +250,7 @@ uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 version = "1.92.2"
 
 [[REPL]]
-deps = ["InteractiveUtils", "Markdown", "Sockets"]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
 uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 
 [[Random]]
@@ -289,8 +290,16 @@ version = "1.1.0"
 deps = ["LinearAlgebra", "SparseArrays"]
 uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
+[[TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+
+[[Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+
 [[Test]]
-deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
+deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [[TimerOutputs]]
@@ -329,7 +338,13 @@ uuid = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
 version = "0.9.3"
 
 [[Zlib_jll]]
-deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
-git-tree-sha1 = "320228915c8debb12cb434c59057290f0834dbf6"
+deps = ["Libdl"]
 uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
-version = "1.2.11+18"
+
+[[nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+
+[[p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
diff --git a/Project.toml b/Project.toml
index b1937c6..7f2d60c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,21 +5,17 @@ version = "0.2.0"
 
 [deps]
 Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
-JSON2 = "2535ab7d-5cd8-5a07-80ac-9b1792aadce3"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
-PackageCompiler = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [compat]
 Conda = "1.4"
-JSON2 = "0.3"
 JuMP = "0.21"
 MathOptInterface = "0.9"
-PackageCompiler = "1"
 PyCall = "1"
 TimerOutputs = "0.5"
 julia = "1"
diff --git a/README.md b/README.md
index 3949252..c141a63 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,21 @@
-# MIPLearn.jl
+<h1 align="center">MIPLearn.jl</h1>
+<p align="center">
+  <a href="https://github.com/ANL-CEEESA/MIPLearn/actions">
+    <img src="https://github.com/ANL-CEEESA/MIPLearn.jl/workflows/Test/badge.svg">
+  </a>
+  <a href="https://doi.org/10.5281/zenodo.4287567">
+    <img src="https://zenodo.org/badge/DOI/10.5281/zenodo.4287567.svg">
+  </a>
+  <a href="https://github.com/ANL-CEEESA/MIPLearn/discussions">
+    <img src="https://img.shields.io/badge/GitHub-Discussions-%23fc4ebc" />
+  </a>
+</p>
 
-This package is a work-in-progress Julia interface for [MIPLearn, an extensible framework for
-Learning-Enhanced Mixed-Integer Optimization](https://github.com/ANL-CEEESA/MIPLearn).
+**MIPLearn** is an extensible open-source framework for solving discrete optimization problems using a combination of Mixed-Integer Linear Programming (MIP) and Machine Learning (ML). See the [main repository](miplearn) for more information. This repository holds an experimental Julia interface for the package.
 
 ## License
 
 Released under the modified BSD license. See `LICENSE` for more details.
+
+
+[miplearn]: https://github.com/ANL-CEEESA/MIPLearn
\ No newline at end of file
diff --git a/src/MIPLearn.jl b/src/MIPLearn.jl
index 8e1cb2c..fa636ba 100644
--- a/src/MIPLearn.jl
+++ b/src/MIPLearn.jl
@@ -1,41 +1,25 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 __precompile__(false)
 module MIPLearn
 
 using PyCall
-miplearn = pyimport("miplearn")
-Instance = miplearn.Instance
-BenchmarkRunner = miplearn.BenchmarkRunner
 
-macro pycall(expr)
-    quote
-        err_msg = nothing
-        result = nothing
-        try
-            result = $(esc(expr))
-        catch err
-            args = err.val.args[1]
-            if (err isa PyCall.PyError) && (args isa String) && startswith(args, "Julia")
-                err_msg = replace(args, r"Stacktrace.*" => "")
-            else
-                rethrow(err)
-            end
-        end
-        if err_msg != nothing
-            error(err_msg)
-        end
-        result
-    end
-end
+export JuMPInstance
+export LearningSolver
+export @feature
+export @category
 
-include("log.jl")
-include("jump_solver.jl")
-include("learning_solver.jl")
-include("instance.jl")
+miplearn = pyimport("miplearn")
 
-export Instance, BenchmarkRunner
+include("utils/log.jl")
+include("utils/pycall.jl")
+include("modeling/jump_instance.jl")
+include("modeling/jump_solver.jl")
+include("modeling/learning_solver.jl")
+include("modeling/macros.jl")
+include("problems/knapsack.jl")
 
 end # module
diff --git a/src/instance.jl b/src/instance.jl
deleted file mode 100644
index 3917ef9..0000000
--- a/src/instance.jl
+++ /dev/null
@@ -1,63 +0,0 @@
-#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
-#  Released under the modified BSD license. See COPYING.md for more details.
-
-using JSON2
-import Base: dump
-
-to_model(instance) =
-    error("not implemented: to_model")
-
-get_instance_features(instance) = [0.0]
-
-get_variable_features(instance, varname) = [0.0]
-
-get_variable_category(instance, varname) = "default"
-
-find_violated_lazy_constraints(instance, model) = []
-
-build_lazy_constraint(instance, model, v) = nothing
-
-macro Instance(klass)
-    quote
-        @pydef mutable struct Wrapper <: Instance
-            function __init__(self, args...; kwargs...)
-                self.data = $(esc(klass))(args...; kwargs...)
-                self.training_data = []
-                self.features = miplearn.Features()
-            end
-                
-            to_model(self) =
-                $(esc(:to_model))(self.data)
-                
-            get_instance_features(self) =
-                $(esc(:get_instance_features))(self.data)
-                
-            get_variable_features(self, varname) =
-                $(esc(:get_variable_features))(self.data, varname)
-
-            get_variable_category(self, varname) =
-                $(esc(:get_variable_category))(self.data, varname)
-
-            find_violated_lazy_constraints(self, model) =
-                find_violated_lazy_constraints(self.data, model)
-            
-            build_lazy_constraint(self, model, v) =
-                build_lazy_constraint(self.data, model, v)
-            
-            load(self) = nothing
-
-            flush(self) = nothing
-        end
-    end
-end
-
-export get_instance_features,
-       get_variable_features,
-       get_variable_category,
-       find_violated_lazy_constraints,
-       build_lazy_constraint,
-       to_model,
-       dump,
-       load!,
-       @Instance
\ No newline at end of file
diff --git a/src/learning_solver.jl b/src/learning_solver.jl
deleted file mode 100644
index d419238..0000000
--- a/src/learning_solver.jl
+++ /dev/null
@@ -1,27 +0,0 @@
-#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
-#  Released under the modified BSD license. See COPYING.md for more details.
-
-struct LearningSolver
-    py::PyCall.PyObject
-end
-
-function LearningSolver(;
-                        optimizer,
-                        kwargs...,
-                       )::LearningSolver
-    py = miplearn.LearningSolver(
-        ;
-        kwargs...,
-        solver=JuMPSolver(optimizer=optimizer),
-    )
-    return LearningSolver(py)
-end
-
-solve!(solver::LearningSolver, instance; kwargs...) =
-    solver.py.solve(instance; kwargs...)
-
-fit!(solver::LearningSolver, instances; kwargs...) =
-    solver.py.fit(instances; kwargs...)
-
-export LearningSolver
\ No newline at end of file
diff --git a/src/modeling/jump_instance.jl b/src/modeling/jump_instance.jl
new file mode 100644
index 0000000..7cbd2c8
--- /dev/null
+++ b/src/modeling/jump_instance.jl
@@ -0,0 +1,52 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+@pydef mutable struct JuMPInstance <: miplearn.Instance
+    function __init__(self, model)
+        init_miplearn_ext(model)
+        features = model.ext[:miplearn][:features]
+        self.model = model
+
+        # Copy training data
+        training_data = []
+        for sample in self.model.ext[:miplearn][:training_samples]
+            pysample = miplearn.TrainingSample()
+            pysample.__dict__ = sample
+            push!(training_data, pysample)
+        end
+        self.training_data = training_data
+
+        # Copy features to data classes
+        self.features = miplearn.Features(
+            instance=miplearn.InstanceFeatures(
+                user_features=PyCall.array2py(
+                    features[:instance][:user_features],
+                ),
+                lazy_constraint_count=0,
+            ),
+            variables=Dict(
+                varname => miplearn.VariableFeatures(
+                    category=vfeatures[:category],
+                    user_features=PyCall.array2py(
+                        vfeatures[:user_features],
+                    ),
+                )
+                for (varname, vfeatures) in features[:variables]
+            ),
+            constraints=Dict(
+                cname => miplearn.ConstraintFeatures(
+                    category=cfeat[:category],
+                    user_features=PyCall.array2py(
+                        cfeat[:user_features],
+                    ),
+                )
+                for (cname, cfeat) in features[:constraints]
+            ),
+        )
+    end
+
+    function to_model(self)
+        return self.model
+    end
+end
diff --git a/src/jump_solver.jl b/src/modeling/jump_solver.jl
similarity index 99%
rename from src/jump_solver.jl
rename to src/modeling/jump_solver.jl
index fcb3590..f80981b 100644
--- a/src/jump_solver.jl
+++ b/src/modeling/jump_solver.jl
@@ -1,5 +1,5 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 using JuMP
diff --git a/src/modeling/learning_solver.jl b/src/modeling/learning_solver.jl
new file mode 100644
index 0000000..f8b6b27
--- /dev/null
+++ b/src/modeling/learning_solver.jl
@@ -0,0 +1,32 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+struct LearningSolver
+    py::PyCall.PyObject
+end
+
+function LearningSolver(
+    ;
+    optimizer,
+)::LearningSolver
+    py = miplearn.LearningSolver(solver=JuMPSolver(optimizer=optimizer))
+    return LearningSolver(py)
+end
+
+function solve!(solver::LearningSolver, model::Model)
+    instance = JuMPInstance(model)
+    mip_stats = solver.py.solve(instance)
+    push!(
+        model.ext[:miplearn][:training_samples],
+        instance.training_data[1].__dict__,
+    )
+    return mip_stats
+end
+
+function fit!(solver::LearningSolver, models::Array{Model})
+    instances = [JuMPInstance(m) for m in models]
+    solver.py.fit(instances)
+end
+
+export LearningSolver
diff --git a/src/modeling/macros.jl b/src/modeling/macros.jl
new file mode 100644
index 0000000..78eedf2
--- /dev/null
+++ b/src/modeling/macros.jl
@@ -0,0 +1,84 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+function init_miplearn_ext(model)::Dict
+    if :miplearn ∉ keys(model.ext)
+        model.ext[:miplearn] = Dict{Symbol,Any}(
+            :features => Dict(
+                :variables => Dict{String,Dict}(),
+                :constraints => Dict{String,Dict}(),
+                :instance => Dict{Symbol,Any}(),
+            ),
+            :training_samples => [],
+        )
+    end
+    return model.ext[:miplearn]
+end
+
+
+function init_miplearn_ext(v::VariableRef)::Dict
+    ext = init_miplearn_ext(v.model)
+    if name(v) ∉ keys(ext[:features][:variables])
+        ext[:features][:variables][name(v)] = Dict{Symbol,Any}()
+    end
+    return ext
+end
+
+
+function init_miplearn_ext(c::ConstraintRef)::Dict
+    ext = init_miplearn_ext(c.model)
+    if name(c) ∉ keys(ext[:features][:constraints])
+        ext[:features][:constraints][name(c)] = Dict{Symbol,Any}()
+    end
+    return ext
+end
+
+
+function set_features!(m::Model, f::Array{Float64})::Nothing
+    ext = init_miplearn_ext(m)
+    ext[:features][:instance][:user_features] = f
+    return
+end
+
+
+function set_features!(v::VariableRef, f::Array{Float64})::Nothing
+    ext = init_miplearn_ext(v)
+    ext[:features][:variables][name(v)][:user_features] = f
+    return
+end
+
+
+function set_category!(v::VariableRef, category::String)::Nothing
+    ext = init_miplearn_ext(v)
+    ext[:features][:variables][name(v)][:category] = category
+    return
+end
+
+
+function set_features!(c::ConstraintRef, f::Array{Float64})::Nothing
+    ext = init_miplearn_ext(c)
+    ext[:features][:constraints][name(c)][:user_features] = f
+    return
+end
+
+
+function set_category!(c::ConstraintRef, category::String)::Nothing
+    ext = init_miplearn_ext(c)
+    ext[:features][:constraints][name(c)][:category] = category
+    return
+end
+
+
+macro feature(obj, features)
+    quote
+        set_features!($(esc(obj)), $(esc(features)))
+    end
+end
+
+
+macro category(obj, category)
+    quote
+        set_category!($(esc(obj)), $(esc(category)))
+    end
+end
diff --git a/src/problems/knapsack.jl b/src/problems/knapsack.jl
new file mode 100644
index 0000000..ac3b6c1
--- /dev/null
+++ b/src/problems/knapsack.jl
@@ -0,0 +1,25 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+using JuMP
+
+function knapsack_model(
+    weights::Array{Float64, 1},
+    prices::Array{Float64, 1},
+    capacity::Float64,
+)
+    model = Model()
+    n = length(weights)
+    @variable(model, x[0:(n-1)], Bin)
+    @objective(model, Max, sum(x[i] * prices[i+1] for i in 0:(n-1)))
+    @constraint(
+        model,
+        eq_capacity,
+        sum(
+            x[i] * weights[i+1]
+            for i in 0:(n-1)
+        ) <= capacity,
+    )
+    return model
+end
diff --git a/src/log.jl b/src/utils/log.jl
similarity index 95%
rename from src/log.jl
rename to src/utils/log.jl
index 005fce6..28d2f2c 100644
--- a/src/log.jl
+++ b/src/utils/log.jl
@@ -1,5 +1,5 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 import Logging: min_enabled_level, shouldlog, handle_message
diff --git a/src/utils/pycall.jl b/src/utils/pycall.jl
new file mode 100644
index 0000000..473f468
--- /dev/null
+++ b/src/utils/pycall.jl
@@ -0,0 +1,24 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+macro pycall(expr)
+    quote
+        err_msg = nothing
+        result = nothing
+        try
+            result = $(esc(expr))
+        catch err
+            args = err.val.args[1]
+            if (err isa PyCall.PyError) && (args isa String) && startswith(args, "Julia")
+                err_msg = replace(args, r"Stacktrace.*" => "")
+            else
+                rethrow(err)
+            end
+        end
+        if err_msg != nothing
+            error(err_msg)
+        end
+        result
+    end
+end
\ No newline at end of file
diff --git a/test/Manifest.toml b/test/Manifest.toml
new file mode 100644
index 0000000..c1717b5
--- /dev/null
+++ b/test/Manifest.toml
@@ -0,0 +1,421 @@
+# This file is machine-generated - editing it directly is not advised
+
+[[ArgTools]]
+uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
+
+[[Artifacts]]
+uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+
+[[Base64]]
+uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
+
+[[BenchmarkTools]]
+deps = ["JSON", "Logging", "Printf", "Statistics", "UUIDs"]
+git-tree-sha1 = "9e62e66db34540a0c919d72172cc2f642ac71260"
+uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+version = "0.5.0"
+
+[[BinaryProvider]]
+deps = ["Libdl", "Logging", "SHA"]
+git-tree-sha1 = "ecdec412a9abc8db54c0efc5548c64dfce072058"
+uuid = "b99e7846-7c00-51b0-8f62-c81ae34c0232"
+version = "0.5.10"
+
+[[Bzip2_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "c3598e525718abcc440f69cc6d5f60dda0a1b61e"
+uuid = "6e34b625-4abd-537c-b88f-471c36dfa7a0"
+version = "1.0.6+5"
+
+[[CEnum]]
+git-tree-sha1 = "215a9aa4a1f23fbd05b92769fdd62559488d70e9"
+uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+version = "0.4.1"
+
+[[Calculus]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "f641eb0a4f00c343bbc32346e1217b86f3ce9dad"
+uuid = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
+version = "0.5.1"
+
+[[Cbc]]
+deps = ["BinaryProvider", "CEnum", "Cbc_jll", "Libdl", "MathOptInterface", "SparseArrays"]
+git-tree-sha1 = "8da071e2f96de95decbd551d54db9ac82c0bf4f6"
+uuid = "9961bab8-2fa3-5c5a-9d89-47fab24efd76"
+version = "0.8.0"
+
+[[Cbc_jll]]
+deps = ["Artifacts", "Cgl_jll", "Clp_jll", "CoinUtils_jll", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "OpenBLAS32_jll", "Osi_jll", "Pkg"]
+git-tree-sha1 = "5d24ea46edebb4f067b180cfc092a45d4d2c48b3"
+uuid = "38041ee0-ae04-5750-a4d2-bb4d0d83d27d"
+version = "2.10.5+2"
+
+[[Cgl_jll]]
+deps = ["Artifacts", "Clp_jll", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "0c41f887cca7e95e7ca1cda5bd6e82e6f43b317d"
+uuid = "3830e938-1dd0-5f3e-8b8e-b3ee43226782"
+version = "0.60.2+6"
+
+[[ChainRulesCore]]
+deps = ["Compat", "LinearAlgebra", "SparseArrays"]
+git-tree-sha1 = "44e9f638aa9ed1ad58885defc568c133010140aa"
+uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+version = "0.9.37"
+
+[[Clp_jll]]
+deps = ["Artifacts", "CoinUtils_jll", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "OpenBLAS32_jll", "Osi_jll", "Pkg"]
+git-tree-sha1 = "d9eca9fa2435959b5542b13409a8ec5f64c947c8"
+uuid = "06985876-5285-5a41-9fcb-8948a742cc53"
+version = "1.17.6+7"
+
+[[CodecBzip2]]
+deps = ["Bzip2_jll", "Libdl", "TranscodingStreams"]
+git-tree-sha1 = "2e62a725210ce3c3c2e1a3080190e7ca491f18d7"
+uuid = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
+version = "0.7.2"
+
+[[CodecZlib]]
+deps = ["TranscodingStreams", "Zlib_jll"]
+git-tree-sha1 = "ded953804d019afa9a3f98981d99b33e3db7b6da"
+uuid = "944b1d66-785c-5afd-91f1-9de20f533193"
+version = "0.7.0"
+
+[[CoinUtils_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "OpenBLAS32_jll", "Pkg"]
+git-tree-sha1 = "5186155a8609b71eae7e104fa2b8fbf6ecd5d9bb"
+uuid = "be027038-0da8-5614-b30d-e42594cb92df"
+version = "2.11.3+4"
+
+[[CommonSubexpressions]]
+deps = ["MacroTools", "Test"]
+git-tree-sha1 = "7b8a93dba8af7e3b42fecabf646260105ac373f7"
+uuid = "bbf7d656-a473-5ed7-a52c-81e309532950"
+version = "0.3.0"
+
+[[Compat]]
+deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
+git-tree-sha1 = "ac4132ad78082518ec2037ae5770b6e796f7f956"
+uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
+version = "3.27.0"
+
+[[CompilerSupportLibraries_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
+
+[[Conda]]
+deps = ["JSON", "VersionParsing"]
+git-tree-sha1 = "6231e40619c15148bcb80aa19d731e629877d762"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.5.1"
+
+[[DataStructures]]
+deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
+git-tree-sha1 = "4437b64df1e0adccc3e5d1adbc3ac741095e4677"
+uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+version = "0.18.9"
+
+[[Dates]]
+deps = ["Printf"]
+uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
+
+[[DelimitedFiles]]
+deps = ["Mmap"]
+uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
+
+[[DiffResults]]
+deps = ["StaticArrays"]
+git-tree-sha1 = "c18e98cba888c6c25d1c3b048e4b3380ca956805"
+uuid = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+version = "1.0.3"
+
+[[DiffRules]]
+deps = ["NaNMath", "Random", "SpecialFunctions"]
+git-tree-sha1 = "214c3fcac57755cfda163d91c58893a8723f93e9"
+uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
+version = "1.0.2"
+
+[[Distributed]]
+deps = ["Random", "Serialization", "Sockets"]
+uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+
+[[Downloads]]
+deps = ["ArgTools", "LibCURL", "NetworkOptions"]
+uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+
+[[ForwardDiff]]
+deps = ["CommonSubexpressions", "DiffResults", "DiffRules", "LinearAlgebra", "NaNMath", "Printf", "Random", "SpecialFunctions", "StaticArrays"]
+git-tree-sha1 = "e2af66012e08966366a43251e1fd421522908be6"
+uuid = "f6369f11-7733-5829-9624-2563aa707210"
+version = "0.10.18"
+
+[[HTTP]]
+deps = ["Base64", "Dates", "IniFile", "MbedTLS", "NetworkOptions", "Sockets", "URIs"]
+git-tree-sha1 = "c9f380c76d8aaa1fa7ea9cf97bddbc0d5b15adc2"
+uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
+version = "0.9.5"
+
+[[IniFile]]
+deps = ["Test"]
+git-tree-sha1 = "098e4d2c533924c921f9f9847274f2ad89e018b8"
+uuid = "83e8ac13-25f8-5344-8a64-a9f2b223428f"
+version = "0.5.0"
+
+[[InteractiveUtils]]
+deps = ["Markdown"]
+uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+
+[[JLLWrappers]]
+git-tree-sha1 = "a431f5f2ca3f4feef3bd7a5e94b8b8d4f2f647a0"
+uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
+version = "1.2.0"
+
+[[JSON]]
+deps = ["Dates", "Mmap", "Parsers", "Unicode"]
+git-tree-sha1 = "81690084b6198a2e1da36fcfda16eeca9f9f24e4"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.21.1"
+
+[[JSON2]]
+deps = ["Dates", "Parsers", "Test"]
+git-tree-sha1 = "66397cc6c08922f98a28ab05a8d3002f9853b129"
+uuid = "2535ab7d-5cd8-5a07-80ac-9b1792aadce3"
+version = "0.3.2"
+
+[[JSONSchema]]
+deps = ["HTTP", "JSON", "ZipFile"]
+git-tree-sha1 = "b84ab8139afde82c7c65ba2b792fe12e01dd7307"
+uuid = "7d188eb4-7ad8-530c-ae41-71a32a6d4692"
+version = "0.3.3"
+
+[[JuMP]]
+deps = ["Calculus", "DataStructures", "ForwardDiff", "JSON", "LinearAlgebra", "MathOptInterface", "MutableArithmetics", "NaNMath", "Random", "SparseArrays", "SpecialFunctions", "Statistics"]
+git-tree-sha1 = "e952f49e2242fa21edcf27bbd6c67041685bee5d"
+uuid = "4076af6c-e467-56ae-b986-b466b2749572"
+version = "0.21.6"
+
+[[LibCURL]]
+deps = ["LibCURL_jll", "MozillaCACerts_jll"]
+uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
+
+[[LibCURL_jll]]
+deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
+uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
+
+[[LibGit2]]
+deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
+uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
+
+[[LibSSH2_jll]]
+deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
+uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
+
+[[Libdl]]
+uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+
+[[LinearAlgebra]]
+deps = ["Libdl"]
+uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+
+[[Logging]]
+uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"
+
+[[MIPLearn]]
+deps = ["Conda", "JSON2", "JuMP", "Logging", "MathOptInterface", "PackageCompiler", "Printf", "PyCall", "TimerOutputs"]
+git-tree-sha1 = "86a0bef09e21ec65ee5becf7d662315bb325b1c6"
+uuid = "2b1277c3-b477-4c49-a15e-7ba350325c68"
+version = "0.1.0"
+
+[[MacroTools]]
+deps = ["Markdown", "Random"]
+git-tree-sha1 = "6a8a2a625ab0dea913aba95c11370589e0239ff0"
+uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+version = "0.5.6"
+
+[[Markdown]]
+deps = ["Base64"]
+uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
+
+[[MathOptInterface]]
+deps = ["BenchmarkTools", "CodecBzip2", "CodecZlib", "JSON", "JSONSchema", "LinearAlgebra", "MutableArithmetics", "OrderedCollections", "SparseArrays", "Test", "Unicode"]
+git-tree-sha1 = "606efe4246da5407d7505265a1ead72467528996"
+uuid = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+version = "0.9.20"
+
+[[MbedTLS]]
+deps = ["Dates", "MbedTLS_jll", "Random", "Sockets"]
+git-tree-sha1 = "1c38e51c3d08ef2278062ebceade0e46cefc96fe"
+uuid = "739be429-bea8-5141-9913-cc70e7f3736d"
+version = "1.0.3"
+
+[[MbedTLS_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
+
+[[Mmap]]
+uuid = "a63ad114-7e13-5084-954f-fe012c677804"
+
+[[MozillaCACerts_jll]]
+uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
+
+[[MutableArithmetics]]
+deps = ["LinearAlgebra", "SparseArrays", "Test"]
+git-tree-sha1 = "ff3aa3e4dbc837f80c2031de2f90125c8b3793f3"
+uuid = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
+version = "0.2.15"
+
+[[NaNMath]]
+git-tree-sha1 = "bfe47e760d60b82b66b61d2d44128b62e3a369fb"
+uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+version = "0.3.5"
+
+[[NetworkOptions]]
+uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
+
+[[OpenBLAS32_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "ba4a8f683303c9082e84afba96f25af3c7fb2436"
+uuid = "656ef2d0-ae68-5445-9ca0-591084a874a2"
+version = "0.3.12+1"
+
+[[OpenSpecFun_jll]]
+deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "9db77584158d0ab52307f8c04f8e7c08ca76b5b3"
+uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
+version = "0.5.3+4"
+
+[[OrderedCollections]]
+git-tree-sha1 = "4fa2ba51070ec13fcc7517db714445b4ab986bdf"
+uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+version = "1.4.0"
+
+[[Osi_jll]]
+deps = ["Artifacts", "CoinUtils_jll", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "OpenBLAS32_jll", "Pkg"]
+git-tree-sha1 = "ef540e28c9b82cb879e33c0885e1bbc9a1e6c571"
+uuid = "7da25872-d9ce-5375-a4d3-7a845f58efdd"
+version = "0.108.5+4"
+
+[[PackageCompiler]]
+deps = ["Libdl", "Pkg", "UUIDs"]
+git-tree-sha1 = "d448727c4b86be81b225b738c88d30334fda6779"
+uuid = "9b87118b-4619-50d2-8e1e-99f35a4d4d9d"
+version = "1.2.5"
+
+[[Parsers]]
+deps = ["Dates"]
+git-tree-sha1 = "c8abc88faa3f7a3950832ac5d6e690881590d6dc"
+uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
+version = "1.1.0"
+
+[[Pkg]]
+deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
+uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+
+[[Printf]]
+deps = ["Unicode"]
+uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+
+[[PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "dd1a970b543bd02efce2984582e996af28cab27f"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.92.2"
+
+[[REPL]]
+deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
+uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
+
+[[Random]]
+deps = ["Serialization"]
+uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+
+[[SHA]]
+uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+
+[[Serialization]]
+uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+
+[[SharedArrays]]
+deps = ["Distributed", "Mmap", "Random", "Serialization"]
+uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
+
+[[Sockets]]
+uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
+
+[[SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[SpecialFunctions]]
+deps = ["ChainRulesCore", "OpenSpecFun_jll"]
+git-tree-sha1 = "5919936c0e92cff40e57d0ddf0ceb667d42e5902"
+uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
+version = "1.3.0"
+
+[[StaticArrays]]
+deps = ["LinearAlgebra", "Random", "Statistics"]
+git-tree-sha1 = "2f01a51c23eed210ff4a1be102c4cc8236b66e5b"
+uuid = "90137ffa-7385-5640-81b9-e52037218182"
+version = "1.1.0"
+
+[[Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[[TOML]]
+deps = ["Dates"]
+uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
+
+[[Tar]]
+deps = ["ArgTools", "SHA"]
+uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
+
+[[Test]]
+deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
+uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[[TimerOutputs]]
+deps = ["Printf"]
+git-tree-sha1 = "32cdbe6cd2d214c25a0b88f985c9e0092877c236"
+uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
+version = "0.5.8"
+
+[[TranscodingStreams]]
+deps = ["Random", "Test"]
+git-tree-sha1 = "7c53c35547de1c5b9d46a4797cf6d8253807108c"
+uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
+version = "0.9.5"
+
+[[URIs]]
+git-tree-sha1 = "7855809b88d7b16e9b029afd17880930626f54a2"
+uuid = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
+version = "1.2.0"
+
+[[UUIDs]]
+deps = ["Random", "SHA"]
+uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+
+[[Unicode]]
+uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
+
+[[VersionParsing]]
+git-tree-sha1 = "80229be1f670524750d905f8fc8148e5a8c4537f"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.2.0"
+
+[[ZipFile]]
+deps = ["Libdl", "Printf", "Zlib_jll"]
+git-tree-sha1 = "c3a5637e27e914a7a445b8d0ad063d701931e9f7"
+uuid = "a5390f91-8eb1-5f08-bee0-b1d1ffed6cea"
+version = "0.9.3"
+
+[[Zlib_jll]]
+deps = ["Libdl"]
+uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
+
+[[nghttp2_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
+
+[[p7zip_jll]]
+deps = ["Artifacts", "Libdl"]
+uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
diff --git a/test/Project.toml b/test/Project.toml
index 9ec49cc..463b9a1 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,6 +1,10 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
 [deps]
+Cbc = "9961bab8-2fa3-5c5a-9d89-47fab24efd76"
 Conda = "8f4d0f93-b110-5947-807f-2305c1781a2d"
-Gurobi = "2e9cd046-0924-5485-92f1-d5272153d98b"
 JSON2 = "2535ab7d-5cd8-5a07-80ac-9b1792aadce3"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MIPLearn = "2b1277c3-b477-4c49-a15e-7ba350325c68"
diff --git a/test/jump_solver_test.jl b/test/jump_solver_test.jl
deleted file mode 100644
index ec5d330..0000000
--- a/test/jump_solver_test.jl
+++ /dev/null
@@ -1,23 +0,0 @@
-#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
-#  Released under the modified BSD license. See COPYING.md for more details.
-
-using Test
-using MIPLearn
-using Gurobi
-using PyCall
-
-miplearn_tests = pyimport("miplearn.solvers.tests")
-
-@testset "JuMPSolver" begin
-    for optimizer in [Gurobi.Optimizer]
-        instance = KnapsackInstance(
-            [23., 26., 20., 18.],
-            [505., 352., 458., 220.],
-            67.0,
-        )
-        model = instance.to_model()
-        solver = JuMPSolver(optimizer=optimizer)
-        miplearn_tests.test_internal_solver(solver, instance, model)
-    end
-end
\ No newline at end of file
diff --git a/test/knapsack.jl b/test/knapsack.jl
deleted file mode 100644
index 13535c2..0000000
--- a/test/knapsack.jl
+++ /dev/null
@@ -1,41 +0,0 @@
-#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
-#  Released under the modified BSD license. See COPYING.md for more details.
-
-import MIPLearn: to_model,
-                 get_instance_features,
-                 get_variable_features
-                 find_violated_lazy_constraints
-using JuMP
-
-struct KnapsackData
-    weights
-    prices
-    capacity
-end
-
-function to_model(data::KnapsackData)
-    model = Model()
-    n = length(data.weights)
-    @variable(model, x[0:(n-1)], Bin)
-    @objective(model, Max, sum(x[i] * data.prices[i+1] for i in 0:(n-1)))
-    @constraint(
-        model,
-        eq_capacity,
-        sum(
-            x[i] * data.weights[i+1]
-            for i in 0:(n-1)
-        ) <= data.capacity,
-    )
-    return model
-end
-
-function get_instance_features(data::KnapsackData)
-    return [0.]
-end
-
-function get_variable_features(data::KnapsackData, var, index)
-    return [0.]
-end
-
-KnapsackInstance = @Instance(KnapsackData)
diff --git a/test/learning_solver_test.jl b/test/learning_solver_test.jl
deleted file mode 100644
index 6091ba0..0000000
--- a/test/learning_solver_test.jl
+++ /dev/null
@@ -1,25 +0,0 @@
-#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
-#  Released under the modified BSD license. See COPYING.md for more details.
-
-using Test
-using MIPLearn
-using Gurobi
-
-
-@testset "LearningSolver" begin
-    for optimizer in [Gurobi.Optimizer]
-        instance = KnapsackInstance(
-            [23., 26., 20., 18.],
-            [505., 352., 458., 220.],
-            67.0,
-        )
-        solver = LearningSolver(
-            optimizer=optimizer,
-            mode="heuristic",
-        )
-        stats = solve!(solver, instance)
-        fit!(solver, [instance])
-        solve!(solver, instance)
-    end
-end
diff --git a/test/modeling/jump_solver_test.jl b/test/modeling/jump_solver_test.jl
new file mode 100644
index 0000000..c6c3fac
--- /dev/null
+++ b/test/modeling/jump_solver_test.jl
@@ -0,0 +1,21 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+using Test
+using MIPLearn
+using Cbc
+using PyCall
+
+miplearn_tests = pyimport("miplearn.solvers.tests")
+
+@testset "JuMPSolver" begin
+    model = MIPLearn.knapsack_model(
+        [23., 26., 20., 18.],
+        [505., 352., 458., 220.],
+        67.0,
+    )
+    instance = JuMPInstance(model)
+    solver = JuMPSolver(optimizer=Cbc.Optimizer)
+    miplearn_tests.test_internal_solver(solver, instance, model)
+end
\ No newline at end of file
diff --git a/test/modeling/learning_solver_test.jl b/test/modeling/learning_solver_test.jl
new file mode 100644
index 0000000..e4ef448
--- /dev/null
+++ b/test/modeling/learning_solver_test.jl
@@ -0,0 +1,75 @@
+#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
+#  Released under the modified BSD license. See COPYING.md for more details.
+
+using JuMP
+using MIPLearn
+using Cbc
+
+@testset "macros" begin
+    weights = [1.0, 2.0, 3.0]
+    prices = [5.0, 6.0, 7.0]
+    capacity = 3.0
+
+    # Create standard JuMP model
+    model = Model()
+    n = length(weights)
+    @variable(model, x[1:n], Bin)
+    @objective(model, Max, sum(x[i] * prices[i] for i in 1:n))
+    @constraint(model, c1, sum(x[i] * weights[i] for i in 1:n) <= capacity)
+
+    # Add machine-learning information
+    @feature(model, [5.0])
+    @feature(c1, [1.0, 2.0, 3.0])
+    @category(c1, "c1")
+    for i in 1:n
+        @feature(x[i], [weights[i]; prices[i]])
+        @category(x[i], "type-$i")
+    end
+ 
+    # Should store variable features
+    @test model.ext[:miplearn][:features][:variables] == Dict(
+        "x[1]" => Dict(
+            :user_features => [1.0, 5.0],
+            :category => "type-1",
+
+        ),
+        "x[2]" => Dict(
+            :user_features => [2.0, 6.0],
+            :category => "type-2",
+        ),
+        "x[3]" => Dict(
+            :user_features => [3.0, 7.0],
+            :category => "type-3",
+        ),
+    )
+
+    # Should store constraint features
+    @test model.ext[:miplearn][:features][:constraints] == Dict(
+        "c1" => Dict(
+            :user_features => [1.0, 2.0, 3.0],
+            :category => "c1",
+        )
+    )
+
+    # Should store instance features
+    @test model.ext[:miplearn][:features][:instance] == Dict(
+        :user_features => [5.0],
+    )
+
+    solver = LearningSolver(optimizer=Cbc.Optimizer)
+    
+    # Should return correct stats
+    stats = solve!(solver, model)
+    @test stats["Lower bound"] == 11.0
+
+    # Should add a sample to the training data
+    @test length(model.ext[:miplearn][:training_samples]) == 1
+    sample = model.ext[:miplearn][:training_samples][1]
+    @test sample["lower_bound"] == 11.0
+    @test sample["solution"]["x[1]"] == 1.0
+
+    fit!(solver, [model])
+
+    solve!(solver, model)
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index 6630ccf..3a564cd 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,5 +1,5 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
-#  Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
 
 using Test
@@ -8,7 +8,6 @@ using MIPLearn
 MIPLearn.setup_logger()
 
 @testset "MIPLearn" begin
-    include("knapsack.jl")
-    include("jump_solver_test.jl")
-    include("learning_solver_test.jl")
+    #include("modeling/jump_solver_test.jl")
+    include("modeling/learning_solver_test.jl")
 end
\ No newline at end of file
diff --git a/test/sysimage.jl b/test/sysimage.jl
deleted file mode 100644
index 93c942f..0000000
--- a/test/sysimage.jl
+++ /dev/null
@@ -1,19 +0,0 @@
-using PackageCompiler
-
-using JSON2
-using CPLEX
-using Gurobi
-using JuMP
-using MathOptInterface
-using PyCall
-using TimerOutputs
-
-pkg = [:JSON2
-       :CPLEX
-       :Gurobi
-       :JuMP
-       :MathOptInterface
-       :PyCall
-       :TimerOutputs]
-
-create_sysimage(pkg, sysimage_path="build/sysimage.so")
\ No newline at end of file